docs: def mbar_log_W_nk(u_kn, N_k, f_k):
      """Calculate the log weight matrix.
  
      Parameters
      ----------
      u_kn : np.ndarray, shape=(n_states, n_samples), dtype='float'
          The reduced potential energies, i.e. -log unnormalized probabilities
      N_k : np.ndarray, shape=(n_states), dtype='int'
          The number of samples in each state
      f_k : np.ndarray, shape=(n_states), dtype='float'
          The reduced free energies of each state
  
      Returns
      -------
      logW_nk : np.ndarray, dtype='float', shape=(n_samples, n_states)
          The normalized log weights.
  
      Notes
      -----
      Equation (9) in JCP MBAR paper.
      """
      u_kn, N_k, f_k = validate_inputs(u_kn, N_k, f_k)
  
      log_denominator_n = logsumexp(f_k - u_kn.T, b=N_k, axis=1)
      logW = f_k - u_kn.T - log_denominator_n[:, np.newaxis]
      return logW