docs: def mbar_gradient(u_kn, N_k, f_k):
      """Gradient of MBAR objective function.
  
      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
      -------
      grad : np.ndarray, dtype=float, shape=(n_states)
          Gradient of mbar_objective
  
      Notes
      -----
      This is equation C6 in the 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)
      log_numerator_k = logsumexp(-log_denominator_n - u_kn, axis=1)
      return -1 * N_k * (1.0 - np.exp(f_k + log_numerator_k))