In 5.x, we were able to pass a jitter value to nuts_sampler_kwargs when calling pm.sample to specify a jitter. nuts_sampler_kwargs is now deprecated in favor of nuts in PyMC 6, but when the same jitter value is passed in nuts, we get this error: TypeError: NUTS._init_() got an unexpected keyword argument ‘jitter’. Is it no longer possible to specify a jitter in PyMC 6?
The current default initialization for NUTS is “jitter+adapt_diag”, so it is being jittered by default. Is there a reason you’d want specific jitter values?
Hi @fonnesbeck! We need to pass some jitter values that works better for the models that we are building. Is there no longer a way to pass jitters now?
In that case you are probably best just setting initial values either by the initval argument of the variable, or as a dictionary of initvals in pm.sample.
We do that as well, and we still want to pass a jitter too
All that the jitter initialization does is add jitter to the initial values, the only difference being that it is applied on the transformed space. So if you just add jitter directly to the initial values you are essentially doing the same thing.
Got it! Thank you, @fonnesbeck !