MBAR's computePerturbedFreeEnergies doesn't play nice with autograd. This does.

autograd_able_compute_perturbed_free_energies.py

from autograd.scipy.misc import logsumexp
from autograd import grad
from autograd import numpy as np
import pymbar

def compute_perturbed_free_energies(mbar, u_ln):
    
    states_with_samples = (mbar.N_k > 0)
    log_q_k = mbar.f_k[states_with_samples] - mbar.u_kn[states_with_samples].T
    log_denominator_n = logsumexp(log_q_k, b=mbar.N_k[states_with_samples], axis=1)
    
    return logsumexp(-u_ln[:len(u_ln)] - log_denominator_n, axis=1)