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

ricardoV94 July 15, 2022, 8:33pm 6

Also, disabling the loglikelihood computation via idata_kwargs if you are not planning to do model comparison can help if your model has many observations

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