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

junpenglao December 12, 2022, 5:44am 10

This is really nice!! How is the speed compare to vmap?

Would you like to send a PR to PyMC?

1 Like

Topic		Replies	Views
Pymc, numpyro GPs and Transform RVs memory behaviour v5 gpu , gaussian_process , sampling	9	896	December 21, 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	513	October 9, 2023
Batch process capability for pymc.sampling_jax.sample_numpyro_nuts() with GPU? v5 modeling	3	583	September 12, 2022
Reduce memory requirements on the GPU when sampling with pm.sampling_jax.sample_numpyro_nuts() v5 gpu	3	1249	March 15, 2023
Pymc/numpyro GPU memory allocation	8	930	June 13, 2022

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

Related topics