Contract of momentum_generator?

[femtomc]

Hi! I've started using BlackJAX recently. I create a "one step" function and pass it over the line to blackjax, but I'm running into some Pytree issues:

I'm not sure what the contract is on momentum_generator? Does this seem like a common error I'm making?

[junpenglao]

momentum_generator input signature is

def momentum_generator(rng_key: PRNGKey, position: PyTree) -> PyTree:

Where the return pytree should be the same structure as the input pytree position. Could you share more how you are calling it?

[femtomc]

Here's how I'm invoking it:

I have a scorer function which takes a Pytree which is split into two pieces (one is differentiable, the other is not).

# The nograd component never changes.
def _logpdf(grad):
    return scorer(grad, nograd)

hmc = blackjax.hmc(
    _logpdf,
    self.step_size,
    self.inverse_mass_matrix,
    self.num_integration_steps,
)

# Pass the grad component into the HMC init.
initial_state = hmc.init(grad)

def step(state, key):
    state, _ = hmc.step(key, state)
    return state, state

sub_keys = jax.random.split(key, self.num_steps)
_, states = jax.lax.scan(step, initial_state, sub_keys)

# ERROR occurs before these lines
final_positions, _ = tree_zipper(states.position, nograd)
return final_positions

[femtomc]

I assume this is a common pattern -- any ideas?

[femtomc]

Ah, I found my bug -- it has to do with jax.scipy.mv_normal in particular -- investigating.