docs: def precondition_u_kn(u_kn, N_k, f_k):
      """Subtract a sample-dependent constant from u_kn to improve precision
  
      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
      -------
      u_kn : np.ndarray, shape=(n_states, n_samples), dtype='float'
          The reduced potential energies, i.e. -log unnormalized probabilities
  
      Notes
      -----
      Returns u_kn - x_n, where x_n is based on the current estimate of f_k.
      Upon subtraction of x_n, the MBAR objective function changes by an
      additive constant, but its derivatives remain unchanged.  We choose
      x_n such that the current objective function value is zero, which
      should give maximum precision in the objective function.
      """
      u_kn, N_k, f_k = validate_inputs(u_kn, N_k, f_k)
      u_kn = u_kn - u_kn.min(0)
      u_kn += (logsumexp(f_k - u_kn.T, b=N_k, axis=1)) - N_k.dot(f_k) / float(N_k.sum())
      return u_kn