{"record":{"id":"0f9ceea7104f8cdc","repo":"jax-ml/jax","slug":"jax-encountered-invalid-prng-key-data-expected-ke-0f9cee","errorCode":null,"errorMessage":"JAX encountered invalid PRNG key data: expected key_data.shape to end with {impl.key_shape}; got shape={key_data.shape} for {impl=}","messagePattern":"JAX encountered invalid PRNG key data: expected key_data\\.shape to end with (.+?); got shape=(.+?) for (.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/prng.py","lineNumber":135,"sourceCode":"\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\n  they do not expose arithmetic operations. They instead expose\n  wrapper methods around the PRNG implementation functions (``split``,\n  ``random_bits``, ``fold_in``).","sourceCodeStart":117,"sourceCodeEnd":153,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/prng.py#L117-L153","documentation":"Every JAX RNG implementation defines a fixed key_shape (e.g. (2,) for threefry2x32, (4,) for rbg). This TypeError means the key_data's trailing dimensions do not match that shape, so the data cannot be interpreted as keys of the requested implementation. It is raised by _check_prng_key_data when key_data.shape[-len(impl.key_shape):] != impl.key_shape.","triggerScenarios":"Wrapping a length-4 uint32 array as a 'threefry2x32' key (expects trailing shape (2,)); wrapping a (2,)-shaped array as an 'rbg' key (expects (4,)); wrapping data with an extra partial dimension such as shape (3,) or (5,).","commonSituations":"Switching the impl argument when migrating from legacy PRNGKey arrays to new-style keys or to rbg; mixing key data from different generator families (threefry vs rbg vs unsafe_rbg) in the same codebase.","solutions":["Match the data to the impl: use 2 uint32 words for threefry2x32, 4 for rbg","Re-derive keys with jax.random.key(seed, impl=impl) instead of transplanting raw data between implementations","Print impl.key_shape for your implementation and reshape/extend the data accordingly"],"exampleFix":"// before\nimpl = jax.random.rbg_prng_impl\nkey = jax.random.random_wrap(jnp.uint32([1, 2]), impl=impl)  # rbg needs 4 words\n\n// after\nkey = jax.random.key(42, impl='rbg')","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\nKEY_SHAPES = {'threefry2x32': (2,), 'rbg': (4,), 'unsafe_rbg': (4,)}\nassert data.shape[-len(KEY_SHAPES[impl_name]):] == KEY_SHAPES[impl_name]","typeGuard":"import jax.numpy as jnp\ndef matches_impl_shape(data, impl) -> bool:\n    ks = impl.key_shape\n    return data.shape[-len(ks):] == ks","tryCatchPattern":null,"preventionTips":["Regenerate keys with jax.random.key(seed, impl=...) when switching implementations","Keep key data paired with its impl name in storage"],"tags":["jax","prng","shape-mismatch","typeerror"],"backgroundTag":"jax-prng-key-validation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}