DOC: pm.sample doesn't use ZarrTrace with numpyro, but docs don't say this

[brendan-m-murphy]

Issue with current documentation:

I tried out the ZarrTrace backend and didn't notice anything being written. A directory is created when I make the ZarrTrace object, but it isn't populated with any of the default group, and doesn't have the trace.

Looking through the code for pm.sample, I noticed that the backend isn't used if you specify e.g. nuts_sampler="numpyro". It would be nice to have this feature with numpyro, but perhaps for now it would be useful to mention in the docs that it only works with PyMC's samplers.

Idea or request for content:

No response

[ricardoV94]

Yeah this doesn't work with external samplers.

It should be possible to work with nutpie, but some piping logic is needed.

I don't think it's easy to work with jax based samplers, at least without taking control over the step routine to do one step at a time

[ricardoV94]

You can however use the regular pymc sampler with jax backend instead of numpyro. Passing compile_kwargs=dict(mode="JAX") to pm.sample.