While this sampler can compile the underlying model to different backends (C, Numba, or JAX) using PyTensor's compilation system via thecompile_kwargs parameter, it maintains Python overhead that can limit performance for large models

One of the masters can correct me, but I think that the python overhead limits performance onsmaller models. For big models, most of the compute time is going to be spent inside thelogp ordlogp function, so doing the python looping will be a lower relative cost.

Also don't forget the poor Torch backend :)

Nutpie is PyMC's cutting-edge performance sampler. Written in Rust, it eliminates Python overhead and provides exceptional performance for continuous models. The Numba backend typically offers the highest performance for most use cases, while the JAX backend excels with very large models and provides GPU acceleration capabilities. Nutpie is particularly well-suited for production workflows where sampling speed is critical.

I'd also mention that Nutpie has the SOTA NUTS adaptation algorithm, so it gets into the typical set much faster and you can get away with many fewer tuning steps as a result. That means even more speedup!

Reply viaReviewNB

torch is just a compile mode, yes?

Yep

Yes@jessegrabowski is correct that performance penalty should be mostly for small models, (ignoring the questions of adaptation, that are not related to the language used, but the algorithm)

Any reason why this model specifically?

Reply viaReviewNB

It was inherited from the previous version of the notebook. No idea. Open to a better (non-simulated) suggestion.

Is this plot needed? You never do anything with the sampling results anyway

Reply viaReviewNB

I think just to give the reader an idea of what the data look like. Not needed; just something I did not remove from the current version of the notebook.

If it's inherited I'm not going to make a fuss, but I'd still consider removing it.

I'm going to replace this example entirely. Looking at a GP model.

Why straight to jax, instead of defaults?

Reply viaReviewNB

You should setxla_force_num_host_devices=8 to make the comparison fair. You're sampling sequentially in this example.

Reply viaReviewNB

I thoughtchain_method="parallel" was the default?

OK, I've addednumpyro.set_host_device_count(8) which seems to do the trick.

I also thought it was the default, but apparently not.

Internally,numpyro.set_host_device_count(8) just doesos.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'. My preference is that we show setting that flag directly, rather than using the numpyro function. For a long time I thought that numpyro function did something special, so I was importing that whole package even when I was sampling with nutpie.

You don't really need a scenario for this though! I pretty much always use numba mode these days.

Would be nice to do the timings on all 3 backends and show the results in a little table.

You can also mention that you can globally request pytensor to compile everything to a specific backend by putting e.g:

import pytensormode = pytensor.compile.mode.get_mode('NUMBA')pytensor.config.mode = mode

At the top of a script/notebook. Then it's not required to pass compile_kwargs, it will default to numba always (including for post-sampling stuff likesample_posterior_predictive,which can be important)

Reply viaReviewNB

I think it would be good to just run all these and make a little table with timings (compile, sample, wall, es/s)

Reply viaReviewNB

Thought about that, but thought it would be clearer to have each model run in its own cell. Let me think about how best to do this.

I also think every model should be run in it's own cell. But we could still collect timings as variables by doingtime.time() before and after the call to sample. There also might be some profiler magic@ricardoV94 or@aseyboldt know to get the compile time vs the sampling time. My proposed method would just lump them together, which will bias against numba for example.

Lumping them together isn't terrible, as it will give an idea of timing in real-world usage. I guess I can increase the number of samples to more heavily weight sampling time.

Would be good to mention that many step samplers compile their own functions that arenot the same as the logp function compiled by the model, and thatcompile_kwargs are not currently propagated to these. So in cases were you're really optimizing for speed, it can be necessary to manually declare step samplers so you can passcompile_kwargs to them.

For the record,BinaryGibbsMetropolis is not such a sampler -- it will respectcompile_kwargs passed topm.sample.But not all will (e.g.Metropolis andSlice)

Reply viaReviewNB