jax-ml/jax · error · TypeError

JAX array with PRNGKey dtype cannot be converted to a NumPy

Error message

JAX array with PRNGKey dtype cannot be converted to a NumPy array. Use jax.random.key_data(arr) if you wish to extract the underlying integer array.

What it means

Typed PRNG keys (dtype key<fry> etc.) carry metadata that NumPy cannot represent, so implicit conversion via np.array(key) or key.__array__ is blocked. JAX directs you to jax.random.key_data(arr), which returns the underlying uint32 bit array, for cases where you truly need raw NumPy data (e.g. serialization).

Source

Thrown at jax/_src/random/prng.py:335

  def pprint(self):
    pp_keys = pp.text('shape = ') + pp.text(str(self.shape))
    pp_impl = pp.text('impl = ') + self._impl.pprint()
    return str(pp.group(
      pp.text('PRNGKeyArray:') +
      pp.nest(2, pp.brk() + pp_keys + pp.brk() + pp_impl)))

  def copy(self):
    out = self.__class__(self._impl, self._base_array.copy())
    out._consumed = self._consumed  # TODO(jakevdp): is this correct?
    return out

  __hash__ = None
  __array_priority__ = 100

  def __array__(self, dtype: np.dtype | None = None, context: None = None,
                copy: bool | None = None) -> np.ndarray:
    del dtype, context, copy
    raise TypeError("JAX array with PRNGKey dtype cannot be converted to a NumPy array."
                    " Use jax.random.key_data(arr) if you wish to extract the underlying"
                    " integer array.")


  # Overwritten immediately below
  @property
  def at(self)                  -> _IndexUpdateHelper: assert False  # pyrefly: ignore[bad-override]
  @property
  def T(self)                   -> PRNGKeyArray: assert False
  def __getitem__(self, _, /)   -> PRNGKeyArray: assert False
  def flatten(self, *_, **__)   -> PRNGKeyArray: assert False
  def ravel(self, *_, **__)     -> PRNGKeyArray: assert False
  def reshape(self, *_, **__)   -> PRNGKeyArray: assert False
  def squeeze(self, *_, **__)   -> PRNGKeyArray: assert False
  def swapaxes(self, *_, **__)  -> PRNGKeyArray: assert False
  def take(self, *_, **__)      -> PRNGKeyArray: assert False
  def transpose(self, *_, **__) -> PRNGKeyArray: assert False

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Call jax.random.key_data(key) to get the uint32 array, then convert to NumPy
2. For round-tripping, store key_data and reconstruct with jax.random.random_wrap(key_data, impl='threefry2x32')
3. Update generic serialization helpers to special-case PRNG keys

Example fix

// before
np.save('key.npy', np.asarray(key))  # TypeError

// after
import jax, jax.numpy as jnp
np.save('key.npy', np.asarray(jax.random.key_data(key)))
# later: key = jax.random.random_wrap(jnp.load('key.npy'), impl='threefry2x32')

import jax, jax.numpy as jnp
def to_numpy_safe(x):
    return np.asarray(jax.random.key_data(x)) if jax.dtypes.issubdtype(x.dtype, jax.dtypes.prng_key) else np.asarray(x)

import jax
def needs_key_data_extraction(x) -> bool:
    return hasattr(x, 'dtype') and jax.dtypes.issubdtype(x.dtype, jax.dtypes.prng_key)

When it happens

Trigger: np.asarray(key) or np.array(key) on a typed key; passing a key to NumPy APIs (np.save, np.stack) that call __array__; mixing keys into lists converted with np.array.

Common situations: Saving/checkpointing code that blindly np.asarray's its inputs; plotting or logging utilities that convert arguments to NumPy; migrating from legacy uint32 PRNGKey arrays where np.array(key) used to work.

Related errors

• The error occurred in the __reduce__ method, which may indic
• len() of unsized object
• No VJP is available
• Duplicate serialization registration for type `{nodetype}`.
• Duplicate serialization registration for serialized_name `{s

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/325140e8f1eb004c. Report an issue: GitHub.