TypeError: JAX encountered invalid PRNG key data: expected key_data.ndim >= 1; got ndim=0

[SaschaFroelich]

I reinstalled JAX on a new machine and now code that ran before throws an error:

from numpyro.infer import MCMC, NUTS, SA
from  jax import random as jran

mcmc = MCMC(kernel, 
          num_warmup=500, 
          num_samples=500,
          num_chains=1, 
          progress_bar=True)

rng_key = jran.PRNGKey(np.random.randint(10_000))
mcmc.run(rng_key, ...)

Now it throws the error

TypeError: JAX encountered invalid PRNG key data: expected key_data.ndim >= 1; got ndim=0

I don't exactly understand what's the problem, since ndim of the rng_key is 1, and it worked with an earlier version of jax. The only difference now is that I switched from jax 0.4.14 with jaxlib 0.4.14. to jax 0.4.18 with jaxlib 0.4.18+cuda12.cudnn89. Does it maybe have something to do with the cuda support?

Package versions
jax 0.4.18
jaxlib 0.4.18+cuda12.cudnn89
numpyro 0.13.2
python 3.10

Error Trace

File "/home/sascha/anaconda3/envs/jax/lib/python3.10/site-packages/spyder_kernels/py3compat.py", line 356, in compat_exec
   exec(code, globals, locals)

 File "/home/sascha/Desktop/vbm_jax/parameter_recovery.py", line 73, in <module>
   mcmc = inf.perform_grouplevel_inference(agent = agent,

 File "/home/sascha/Desktop/vbm_jax/inference_numpyro.py", line 156, in perform_grouplevel_inference
   mcmc.run(rng_key,

 File "/home/sascha/anaconda3/envs/jax/lib/python3.10/site-packages/numpyro/infer/mcmc.py", line 634, in run
   states_flat, last_state = partial_map_fn(map_args)

 File "/home/sascha/anaconda3/envs/jax/lib/python3.10/site-packages/numpyro/infer/mcmc.py", line 416, in _single_chain_mcmc
   new_init_state = self.sampler.init(

 File "/home/sascha/anaconda3/envs/jax/lib/python3.10/site-packages/numpyro/infer/hmc.py", line 711, in init
   vmap(random.split)(rng_key), 0, 1

 File "/home/sascha/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
   return fun(*args, **kwargs)

 File "/home/sascha/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/api.py", line 1682, in vmap_f
   out_flat = batching.batch(

 File "/home/sascha/anaconda3/envs/jax/lib/python3.10/site-packages/jax/linear_util.py", line 167, in call_wrapped
   ans = self.f(*args, **dict(self.params, **kwargs))

 File "/home/sascha/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/random.py", line 212, in split
   key, wrapped = _check_prng_key(key)

 File "/home/sascha/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/random.py", line 74, in _check_prng_key
   return prng.random_wrap(key, impl=default_prng_impl()), True

 File "/home/sascha/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/prng.py", line 753, in random_wrap
   _check_prng_key_data(impl, base_arr)

 File "/home/sascha/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/prng.py", line 103, in _check_prng_key_data
   raise TypeError("JAX encountered invalid PRNG key data: expected "

TypeError: JAX encountered invalid PRNG key data: expected key_data.ndim >= 1; got ndim=0

[jakevdp]

Thanks for the question - can you share the full error traceback? There's not enough information in your question to know what might be going on. Also, what version of numpyro are you using?

[SaschaFroelich]

Sure, please see my edited post. I'm using numpyro 0.13.2.

[jakevdp]

I think you're using an older JAX version than you report. Your traceback says that line 103 of jax/_src/prng.py is within the _check_prng_key_data function, and contains raise TypeError. In JAX v0.4.18, that's not what this line contains: https://github.com/google/jax/blob/jaxlib-v0.4.18/jax/_src/prng.py#L103. From the contents of this line, it looks like you're using JAX version 0.4.7 or older (see https://github.com/google/jax/blob/jaxlib-v0.4.7/jax/_src/prng.py#L103).

I suspect the error is caused by using an old JAX version that is incompatible with the recent numpyro release: numpyro 0.13.2 specifies that it requires JAX v0.4.14 or newer (https://github.com/pyro-ppl/numpyro/blob/0.13.2/setup.py#L12). I'd suggest you update JAX, and take steps to ensure you're using the package versions that you think you are (for example, you may have multiple python executables in your system with different packages installed in each).

[SaschaFroelich]

You're right, I was relying on my anaconda, which tells me I have jax 0.4.18, but when I execute jax.__version__ in python, it says I have jax 0.3.25. Thank you.