Thanks for your reply! I guess you refer to Pass user-provided NUTS kwargs to Numpyro by jhrcook · Pull Request #6021 · pymc-devs/pymc · GitHub
I looked over it and as I understood it, it only affects ’adapt_step_size’, ’adapt_mass_matrix’ and ’dense_mass’, but I’ll try it out.