{"record":{"id":"325140e8f1eb004c","repo":"jax-ml/jax","slug":"jax-array-with-prngkey-dtype-cannot-be-converted-t","errorCode":null,"errorMessage":"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.","messagePattern":"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\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/prng.py","lineNumber":335,"sourceCode":"  def pprint(self):\n    pp_keys = pp.text('shape = ') + pp.text(str(self.shape))\n    pp_impl = pp.text('impl = ') + self._impl.pprint()\n    return str(pp.group(\n      pp.text('PRNGKeyArray:') +\n      pp.nest(2, pp.brk() + pp_keys + pp.brk() + pp_impl)))\n\n  def copy(self):\n    out = self.__class__(self._impl, self._base_array.copy())\n    out._consumed = self._consumed  # TODO(jakevdp): is this correct?\n    return out\n\n  __hash__ = None\n  __array_priority__ = 100\n\n  def __array__(self, dtype: np.dtype | None = None, context: None = None,\n                copy: bool | None = None) -> np.ndarray:\n    del dtype, context, copy\n    raise TypeError(\"JAX array with PRNGKey dtype cannot be converted to a NumPy array.\"\n                    \" Use jax.random.key_data(arr) if you wish to extract the underlying\"\n                    \" integer array.\")\n\n\n  # Overwritten immediately below\n  @property\n  def at(self)                  -> _IndexUpdateHelper: assert False  # pyrefly: ignore[bad-override]\n  @property\n  def T(self)                   -> PRNGKeyArray: assert False\n  def __getitem__(self, _, /)   -> PRNGKeyArray: assert False\n  def flatten(self, *_, **__)   -> PRNGKeyArray: assert False\n  def ravel(self, *_, **__)     -> PRNGKeyArray: assert False\n  def reshape(self, *_, **__)   -> PRNGKeyArray: assert False\n  def squeeze(self, *_, **__)   -> PRNGKeyArray: assert False\n  def swapaxes(self, *_, **__)  -> PRNGKeyArray: assert False\n  def take(self, *_, **__)      -> PRNGKeyArray: assert False\n  def transpose(self, *_, **__) -> PRNGKeyArray: assert False\n","sourceCodeStart":317,"sourceCodeEnd":353,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/prng.py#L317-L353","documentation":"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).","triggerScenarios":"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.","commonSituations":"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.","solutions":["Call jax.random.key_data(key) to get the uint32 array, then convert to NumPy","For round-tripping, store key_data and reconstruct with jax.random.random_wrap(key_data, impl='threefry2x32')","Update generic serialization helpers to special-case PRNG keys"],"exampleFix":"// before\nnp.save('key.npy', np.asarray(key))  # TypeError\n\n// after\nimport jax, jax.numpy as jnp\nnp.save('key.npy', np.asarray(jax.random.key_data(key)))\n# later: key = jax.random.random_wrap(jnp.load('key.npy'), impl='threefry2x32')","handlingStrategy":"fallback","validationCode":"import jax, jax.numpy as jnp\ndef to_numpy_safe(x):\n    return np.asarray(jax.random.key_data(x)) if jax.dtypes.issubdtype(x.dtype, jax.dtypes.prng_key) else np.asarray(x)","typeGuard":"import jax\ndef needs_key_data_extraction(x) -> bool:\n    return hasattr(x, 'dtype') and jax.dtypes.issubdtype(x.dtype, jax.dtypes.prng_key)","tryCatchPattern":null,"preventionTips":["Special-case prng_key dtypes in serialization helpers","Store key_data (uint32) plus impl name for round-tripping"],"tags":["jax","prng","numpy","serialization"],"backgroundTag":"jax-prng-key-validation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}