-
Notifications
You must be signed in to change notification settings - Fork 2.2k
Open
Milestone
Description
Description
With PyMC 6.0 and PyTensor 3.0 we have numba by default, which nutpie loves.
Let's change the default sampler to be nutpie. This can happen when the following conditions apply:
- All variables are continuous / differentiable
- No sampler gobbled up other variables (BART, custom assignment from the user - unless all in a pm.Nuts step anyway
- compile_kwargs doesn't have a specific mode that is not JAX or Numba Linker
pytensor.config.mode|linerdon't imply a non-JAX/Numba linker
Some changes we want to see on nutpie:
- Use our progressbar (or something that looks pretty much identical)
- Don't store unconsrained variables in the posterior group (either exclude, or store in the separate group that arviz suggests)
- Don't freeze the model by default (there might already be an argument). Our users decide when the model is frozen or not
Changes on PyMC:
- Change the default tune to whatever nutpie is happy about (maybe let it choose)
- Allow the non-blocking object that nutpie provides be returned from pm.sample (optionally), see Non-blocking sampling #7929
twiecki