I am able to finish the sampling process on a large model using the JAX backend with a GPU, but it fails during the “transforming variables” step because it requires too much RAM. I set postprocessing_backend="cpu" which has a lot more RAM than my GPU, but it still runs out. Looking at the code in …

in fact this seems to work: from jax.experimental.maps import SerialLoop, xmap num_chunks = 10 # must be a multiple of number of samples mapper = xmap(jax_fn, in_axes=['chain', 'samples', ...], out_axes=["chain", 'samples', ...], axis_resources={"samples…

Out of memory when "transforming variables" in Numpyro & JAX

Questions v5

wnorcbrown December 12, 2022, 10:02am 11

As far as I could tell the speed was the same.

Sure

1 Like

Pymc, numpyro GPs and Transform RVs memory behaviour

Topic		Replies	Views
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
Reduce memory requirements on the GPU when sampling with pm.sampling_jax.sample_numpyro_nuts() v5 gpu	3	1005	March 15, 2023
Batch process capability for pymc.sampling_jax.sample_numpyro_nuts() with GPU? v5 modeling	3	507	September 12, 2022
Pymc, numpyro GPs and Transform RVs memory behaviour v5 gpu , gaussian_process , sampling	9	823	December 21, 2022
Out of Memory when using pm.sampling.jax.sample_blackjax_nuts v5 gpu , modeling , jax	2	828	March 23, 2023

Out of memory when "transforming variables" in Numpyro & JAX

Related topics