{"record":{"id":"65b642bc7e1f81b8","repo":"jax-ml/jax","slug":"jax-encountered-invalid-prng-key-data-expected-ke","errorCode":null,"errorMessage":"JAX encountered invalid PRNG key data: expected key_data to have ndim, shape, and dtype attributes. Got {key_data}","messagePattern":"JAX encountered invalid PRNG key data: expected key_data to have ndim, shape, and dtype attributes\\. Got (.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/prng.py","lineNumber":129,"sourceCode":"            pp.nest(2, pp.group(pp.brk() + pp.join(pp.brk(), [\n              pp.text(f\"{k} = {v}\") for k, v in self._asdict().items()\n            ]))))\n\n\nprngs: dict[str, PRNGImpl] = {}\n\ndef register_prng(impl: PRNGImpl):\n  if impl.name in prngs:\n    raise ValueError(f'PRNG with name {impl.name} already registered: {impl}')\n  prngs[impl.name] = impl\n\n\n# -- PRNG key arrays\n\ndef _check_prng_key_data(impl, key_data: typing.Array):\n  ndim = len(impl.key_shape)\n  if not all(hasattr(key_data, attr) for attr in ['ndim', 'shape', 'dtype']):\n    raise TypeError(\"JAX encountered invalid PRNG key data: expected key_data \"\n                    f\"to have ndim, shape, and dtype attributes. Got {key_data}\")\n  if key_data.ndim < 1:\n    raise TypeError(\"JAX encountered invalid PRNG key data: expected \"\n                    f\"key_data.ndim >= 1; got ndim={key_data.ndim}\")\n  if key_data.shape[-ndim:] != impl.key_shape:\n    raise TypeError(\"JAX encountered invalid PRNG key data: expected key_data.shape to \"\n                    f\"end with {impl.key_shape}; got shape={key_data.shape} for {impl=}\")\n  if key_data.dtype not in [np.uint32, float0]:\n    raise TypeError(\"JAX encountered invalid PRNG key data: expected key_data.dtype = uint32; \"\n                    f\"got dtype={key_data.dtype}\")\n\n\nclass PRNGKeyArray(Array):\n  \"\"\"An array of PRNG keys backed by an RNG implementation.\n\n  This class lifts the definition of a PRNG, provided in the form of a\n  ``PRNGImpl``, into an array-like pytree class. Instances of this\n  class behave like an array whose base elements are keys, hiding the","sourceCodeStart":111,"sourceCodeEnd":147,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/prng.py#L111-L147","documentation":"JAX's typed PRNG key arrays (PRNGKeyArray) wrap an underlying 'key_data' array whose implementation is validated at construction. This TypeError means the object passed as key_data is not array-like from JAX's perspective: it lacks basic ndim/shape/dtype attributes (e.g. a plain Python list, an int, or None was passed where a key data array was expected). JAX throws it as a fail-fast check inside _check_prng_key_data, which runs from PRNGKeyArray.__init__ and jax.random.random_wrap.","triggerScenarios":"Constructing a PRNGKeyArray/PRNGImpl with a non-array key_data (e.g. PRNGKeyArray(impl, [1,2,3])), or calling jax.random.random_wrap with a Python list/tuple/scalar instead of a uint32 array; also passing objects of array-like libraries that don't expose the trio of attributes.","commonSituations":"Migrating legacy code that stored raw key data as Python lists; deserializing keys from JSON/config and passing the decoded list straight back; wrapping third-party array objects that lack JAX's array API surface.","solutions":["Convert the data to a JAX array first: jax.random.random_wrap(jnp.asarray(data, dtype=jnp.uint32), impl=...)","If starting from a seed, create a key with jax.random.key(seed) or jax.random.PRNGKey(seed) instead of hand-building key data","Ensure the wrapped data is uint32 with trailing shape matching the impl's key_shape (e.g. (2,) for threefry2x32)"],"exampleFix":"// before\nraw = [123, 456]\nkey = jax.random.random_wrap(raw, impl='threefry2x32')\n\n// after\nimport jax.numpy as jnp\nkey = jax.random.random_wrap(jnp.asarray(raw, dtype=jnp.uint32), impl='threefry2x32')","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\ndef is_valid_key_data(x):\n    return (all(hasattr(x, a) for a in ('ndim', 'shape', 'dtype'))\n            and x.ndim >= 1\n            and x.shape[-2:] == (2,)\n            and x.dtype == jnp.uint32)","typeGuard":"import jax.numpy as jnp\nfrom typing import Any\ndef is_wrappable_key_data(x: Any) -> bool:\n    return (all(hasattr(x, a) for a in ('ndim', 'shape', 'dtype'))\n            and x.ndim >= 1 and x.shape[-2:] == (2,)\n            and x.dtype == jnp.uint32)","tryCatchPattern":null,"preventionTips":["Always create keys via jax.random.key(seed) rather than wrapping raw data","Convert inputs with jnp.asarray(..., dtype=jnp.uint32) before random_wrap"],"tags":["jax","prng","typeerror","input-validation"],"backgroundTag":"jax-prng-key-validation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}