While developing GP models with numpyro to sample via the GPU I’ve noticed somewhat strange memory behavior. After pymc.sampling.jax.sample_numpyro_nuts finishes running (very fast) it enters a ‘Transforming RVs’ phase witch triggers a massive memory spike. I’m not sufficiently familiar with the in…

Pymc, numpyro GPs and Transform RVs memory behaviour

ricardoV94 December 14, 2022, 5:56am 3

Topic		Replies	Views
Batch process capability for pymc.sampling_jax.sample_numpyro_nuts() with GPU? v5 modeling	3	507	September 12, 2022
Out of memory when "transforming variables" in Numpyro & JAX v5 jax	10	1190	December 12, 2022
Out of Memory during Transforming Variables Step using Numpyro. Increasing postprocessing_chunks doesn't seems to help. What should I do? v5 modeling	4	425	October 9, 2023
Pymc/numpyro GPU memory allocation	8	794	June 13, 2022
Reduce memory requirements on the GPU when sampling with pm.sampling_jax.sample_numpyro_nuts() v5 gpu	3	1005	March 15, 2023