{"record":{"id":"042dd6db25f6c5f7","repo":"jax-ml/jax","slug":"jax-encountered-invalid-prng-key-data-expected-ke-042dd6","errorCode":null,"errorMessage":"JAX encountered invalid PRNG key data: expected key_data.ndim >= 1; got ndim={key_data.ndim}","messagePattern":"JAX encountered invalid PRNG key data: expected key_data\\.ndim >= 1; got ndim=(.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/prng.py","lineNumber":132,"sourceCode":"\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\n  fact that keys are typically arrays (of ``uint32`` dtype) themselves.\n\n  PRNGKeyArrays are also restricted relative to JAX arrays in that","sourceCodeStart":114,"sourceCodeEnd":150,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/prng.py#L114-L150","documentation":"JAX requires PRNG key_data to be at least 1-dimensional because the trailing dimensions encode the key shape for the RNG implementation. This TypeError fires when a 0-d (scalar) array is passed as key_data, which cannot hold the two uint32 words a threefry key needs. It is raised by _check_prng_key_data during PRNGKeyArray construction or jax.random.random_wrap.","triggerScenarios":"Passing a scalar jnp.uint32 value to jax.random.random_wrap; squeezing/broadcasting key data down to 0-d before wrapping; reshaping key arrays with reshape(()) accidentally.","commonSituations":"Data-processing pipelines that flatten or squeeze arrays generically before passing them on; refactoring code where a key was previously a length-2 vector but became a scalar after indexing.","solutions":["Pass at least a 1-d array whose trailing dims match the impl key_shape: jax.random.random_wrap(jnp.uint32([lo, hi]), impl=...)","If the scalar is a seed, use jax.random.key(seed) instead of wrapping","Check intermediate reshape/squeeze calls that may have collapsed key data to 0-d"],"exampleFix":"// before\nkey = jax.random.random_wrap(jnp.uint32(42), impl='threefry2x32')\n\n// after\nkey = jax.random.key(42)","handlingStrategy":"validation","validationCode":"assert jnp.asarray(data).ndim >= 1, 'key_data must be at least 1-d'","typeGuard":"import jax.numpy as jnp\ndef has_key_ndim(x) -> bool:\n    return hasattr(x, 'ndim') and x.ndim >= 1","tryCatchPattern":null,"preventionTips":["Never pass scalars to random_wrap; use jax.random.key for seeds","Audit squeeze()/reshape(()) calls on key data"],"tags":["jax","prng","ndim","typeerror"],"backgroundTag":"jax-prng-key-validation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}