Execution of replica 0 failed: INVALID_ARGUMENT

[jaxengodfrey]

Hello,

I am using GPU enabled numpyro and jax. I am getting the following error after my chain finishes sampling but before numpyro.infer.MCMC.run() has finished compiling the results:

    mcmc.run(rng_key_, pedict, injdict, constants['nObs'], constants['obs_time'], constants['total_inj'], mass_models, mag_model, tilt_model, z_model, parsargs.mmin, parsargs.mmax, nspline_dict, param_names, rngkey=catkey)
  File "/home/jaxeng/.conda/envs/gwinferno_jaxns/lib/python3.12/site-packages/numpyro/infer/mcmc.py", line 682, in run
    states_flat, last_state = partial_map_fn(map_args)
                              ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jaxeng/.conda/envs/gwinferno_jaxns/lib/python3.12/site-packages/numpyro/infer/mcmc.py", line 500, in _single_chain_mcmc
    states[self._sample_field] = postprocess_fn(states[self._sample_field])
                                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
jaxlib.xla_extension.XlaRuntimeError: INVALID_ARGUMENT: Can't pack device memory arguments array of size 1255 which is larger than the maximum supported size of 1024

I only get this error with large warmup and sample sizes, ~100k. I think it started occurring when I started saving a large number of deterministic variables during sampling. I have done this in the past with older versions of numpyro/jax without issue, same sample sizes, deterministic variables, GPU, etc.

I'm using the NUTS kernel and an 80GB Nvidia A100 GPU.

Because I don't know much about how numpyro uses jax under the hood, I'm not sure how I could isolate this issue within jax itself to make troubleshooting easier. Any ideas/suggestions would be appreciated!