PYMC 5 significant speedup of default sampler (Pytensor)

[zweli]

I installed PYMC 5 this morning (which uses Pytensor as the backend). My machine is an M1 Max with 64 GB memory. I decided to test the default sampler vs Jax-Numpyro and Jax-Blackjax. Prior to PYMC5, the default sampler would always be much slower than Jax.

I loaded up @AustinRochford NBA Foul IRT example at link.

I don’t know what dark magic the PYMC devs are doing (maybe the porting to Numba?)/is this downstream from Aesara but the default sampler was better than both JAX-based samplers (Pytensor=3min20s; Numpyro=3min58s; and BlackJax=5min15s).
I had to rerun a couple of times to confirm because I was initially incredulous. But Pytensor was faster, at least for this model. I have attached an image of the sampling with the different backends.

I will test with other models but if this holds, that would be very impressive.

[zweli]

I am also running on Python 3.11 – so not sure if that is having an effect too.

[michaelosthege]

This is most probably an effect of Python 3,11,
While we’re making small improvements all the time, I’m not aware of anything between v4 and v5 that would explain such a big jump.

I just ran one of my models with PyMC 4.2.2 under Python 3.10.8 vs. PyMC v5.0.0 under Python 3.9.15 and the v5 run was only about 5 % faster.

But the fact that you ran something with Python 3.11 is really good news!
Because the PyTensor GitHub Actions test only 3.7 and 3.9 and the PyMC tests run with Python 3.8, 3.9 and 3.10.

Would you like to have a shot at upgrading the GH actions to Python 3.11 ?
This could save a lot on CI resources!

[zweli]

I don’t know how to do that. If you can walk me through, I can try.

[michaelosthege]

Sorry, I didn’t check my discourse notifications.

Looks like Ricardo opened a WIP PR to test PyTensor with Python 3.11, but it needs someone to look into it: Add support for Python 3.11 and drop support for 3.7 by ricardoV94 · Pull Request #198 · pymc-devs/pytensor · GitHub