hi, I am trying to run a large hierarchical model with pymc + jax on GPU but I keep running into out-of-memory issues. I tried various solutions that I found online but couldn’t really solve the problem so far. Is there anyway to go around this, for example by using batches? I’m very new to pymc a…

Batch process capability for pymc.sampling_jax.sample_numpyro_nuts() with GPU?

ricardoV94 September 12, 2022, 7:29pm 4

That’s more of a JAX/Numpyro question so you may have more luck asking in their forums.

Topic		Replies	Views
Out of Memory when using pm.sampling.jax.sample_blackjax_nuts v5 gpu , modeling , jax	2	917	March 23, 2023
Larger dataset linear regression GPU memory problems with NUTS for larger N and D (batch training?) Questions theano	2	1868	October 3, 2017
Is it possible to speed up PyMC sampling? version agnostic	3	3384	May 20, 2022
GPU is running the model much slower than the CPU v5	7	1544	August 16, 2023
Pymc3 / pymc v4 GPU example for a basic hierarchical model with NUTS sampling version agnostic gpu	1	1273	April 6, 2022