docs: def self_consistent_update(u_kn, N_k, f_k):
      """Return an improved guess for the dimensionless free energies
  
      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
      -------
      f_k : np.ndarray, shape=(n_states), dtype='float'
          Updated estimate of f_k
  
      Notes
      -----
      Equation C3 in MBAR JCP paper.
      """
  
      u_kn, N_k, f_k = validate_inputs(u_kn, N_k, f_k)
      
      states_with_samples = (N_k > 0)
  
      # Only the states with samples can contribute to the denominator term.
      log_denominator_n = logsumexp(f_k[states_with_samples] - u_kn[states_with_samples].T, b=N_k[states_with_samples], axis=1)
      
      # All states can contribute to the numerator term.
      return -1. * logsumexp(-log_denominator_n - u_kn, axis=1)