How can I correctly define my custom gradient and sensitivities for use with NUTS?

[dynova]

I have an ODE model with a model-constrained custom gradient and a custom method to calculate adjoint sensitivities for use with NUTS.

I would like to know how I can properly feed in the gradient and sensitivities into the NUTS sampler.

I went through the documentation here: https://blackjax-devs.github.io/blackjax/examples/howto_custom_gradients.html

But am still confused as to whether I'm doing it correctly or not.

In contrast to the example in the documentation, my sensitivities have no closed-form analytical expression.

Thanks in advance for any help.

The relevant excerpt of my code is below:

@partial(custom_jvp)
def f_with_gradient(parsub):
    par = jnp.array([t0, tdelta, tdelta2, tdelta3, b, lamb, fP, lamb2, fP2, lamb3, fP3, fDI, r])
    par = par.at[freeParameterInd].set(parsub)
    return gradient(par)

def g(parsub):
    par = jnp.array([t0, tdelta, tdelta2, tdelta3, b, lamb, fP, lamb2, fP2, lamb3, fP3, fDI, r])
    par = par.at[freeParameterInd].set(parsub)
    return adjointSensitivities(par)

@f_with_gradient.defjvp
def f_jac_vec_prod(primals, tangents):
    x, = primals
    x_dot, = tangents
    return gradient(par), g(parsub) * x_dot

def inference_loop(rng_key, kernel, initial_state, num_samples):
    @jax.jit
    def one_step(state, rng_key):
        state, _ = kernel(rng_key, state)
        return state, state

    keys = jax.random.split(rng_key, num_samples)
    _, states = jax.lax.scan(one_step, initial_state, keys)
    
    return states.position, states.potential_energy, states.potential_energy_grad

def logdensity_fn(parsub):
    logdensity = f_with_gradient(parsub)
    return logdensity

step_size = 1e-3
inverse_mass_matrix = jnp.atleast_1d(jnp.load('diffMatrix.npy'))
nuts = blackjax.nuts(logdensity_fn, step_size, inverse_mass_matrix)
freeParameterInd = jnp.array([0])
parsub = jnp.atleast_1d(par[freeParameterInd])
initial_position = parsub
initial_state = nuts.init(initial_position)
rng_key = jax.random.PRNGKey(0)
step = nuts.step
n_samples = 2
parLog, logLLog, adjLog = inference_loop(rng_key, step, initial_state, n_samples)

[johannahaffner]

Hi @dynova,

did you figure this out?

I can't quite follow all of your example, I'm trying to because I want to do something similar.