Skip to content

Use nutpie by default (*conditions apply) #8079

@ricardoV94

Description

@ricardoV94

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|liner don'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

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions