pymc on HPC?

[fbarfi]

I am using pymc 4.0 on HPC server. It works fine but much slower than my laptop. Can one use a GPU backend? Thanks for any feedback.

[ferrine]

We have pm.sampling_jax so you can sampler with numpyro or blackjax using jax xla backend

[fbarfi]

Thanks. Ferrine. Yes pm.sampling_jax works wonderfully fast on the HPC. One problem I have is that ODE class does not yet possess an Op that is recognizable by JAX (within pymc or aesara) and I need it in the code that I am running (a sort of Lotka-Volterra problem).

[twiecki]

You could try to use the JAX ODE solver and wrap that into an aesara op.

[fbarfi]

thanks, the idea looks indeed promising.
on the other hand, I just found jax.experimental.ode used in numpyro in an example solving a Lotka-Volterra problem. I'll see which one is doable more effectively for me.

[fbarfi]

is this example from the pymc site a good starting point?

How to wrap a JAX function for use in PyMC

[twiecki]

Yes, exactly.

[twiecki]

Figuring this out would be super useful for other people as well and make a great example NB.

[fbarfi]

Thanks. I'll give it a shot and see where it takes me.

[michaelosthege]

@fbarfi did you try sunode already? Just this morning we did a 0.3.0-alpha release with PyMC v4+ support.
The performance should be much better than pm.ode.DifferentialEquation and need less tinkering than the JAX stuff.

[fbarfi]

thank you @michaelosthege for the suggestion. I tried before to install sunode on my M1 Max but could not do it. On the other hand I made some progress on following @twiecki suggestion. I wrote an aesara Op using scipy odeint (and another one using the jax odeint) and it works great with pm.sample(). I am having some errors when I try to jaxify (for using sampling_jax) which I am trying to figure out -- the problem currently has to do with messages like

TypeError: Shapes must be 1D sequences of concrete values of integer type, got (50, Traced<ShapedArray(int64[])>with<DynamicJaxprTrace(level=0/1)>).

or

The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ShapedArray(float64[4])>with<DynamicJaxprTrace(level=0/1)>

I am learning how to solve these issues.

[fbarfi]

I should have added that using odeint from Jax in my aesara Op is much faster than using odeint from scipy even without sampling_jax.

[fbarfi]

This morning I discovered that my problem (despite of the fact that the error messages were phrased in terms of Jax traces) originated in using MvNormal for the likelihood in my model. Keeping everything the same and replacing MvNormal with Normal made the problem disappear. pm.sampling_jax is working fine. However, it is much slower than pm.sample. best.