Retrieving Hyperparameters after Tuning

Hello everybody,
Is there a way to retrieve the tuned hyperparameters for the base HMC sampler (stepsize, Mass-matrix) and for the NUTS sampler (Mass matrix) after the sampling?
I greatly appreciate your help!
Thanks,
Kevin

Welcome!

Have you checked out the pm.init_nuts() routine?

You can access the mass matrix if you manually create the step method, and don’t use the parallel sampler by setting cores=1:

import pymc as pm

with pm.Model() as model:
    pm.Normal("alpha")

    step = pm.NUTS()
    pm.sample(step=step, cores=1, chains=1)

step.potential._stds

If you sample using nutpie, you can also store the current mass matrix in the sampler stats:

compiled = nutpie.compile_pymc_model(model)
tr = nutpie.sample(compiled, store_mass_matrix=True)
tr.warmup_sample_stats.mass_matrix_inv

The stepsize is always stored in the sample_stats.

1 Like

Thank you very much!