{"record":{"id":"83d9b1b41062f341","repo":"jax-ml/jax","slug":"jax-encountered-invalid-prng-key-data-expected-ke-83d9b1","errorCode":null,"errorMessage":"JAX encountered invalid PRNG key data: expected key_data.dtype = uint32; got dtype={key_data.dtype}","messagePattern":"JAX encountered invalid PRNG key data: expected key_data\\.dtype = uint32; got dtype=(.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/prng.py","lineNumber":138,"sourceCode":"    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``).\n  \"\"\"\n  # TODO(jakevdp): potentially add tolist(), tobytes(),\n  #    device_buffer, device_buffers, __cuda_interface__()","sourceCodeStart":120,"sourceCodeEnd":156,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/prng.py#L120-L156","documentation":"JAX PRNG key data must be stored as uint32 (or the internal float0 sentinel). This TypeError fires when key_data has any other dtype (int64, float32, etc.), because the RNG bit-manipulation kernels only operate on 32-bit unsigned words. The check lives in _check_prng_key_data and runs from PRNGKeyArray.__init__ and jax.random.random_wrap.","triggerScenarios":"Calling jax.random.random_wrap on an int64 array (common when x64 mode is enabled), a float array, or a numpy array of dtype int32; converting seeds with jnp.array(seed) (default int32/int64) and wrapping directly.","commonSituations":"JAX_ENABLE_X64=1 environments where asarray defaults to int64; loading key data from files where it was stored as int64 or float; numpy interop where np.array([1,2]) yields int64 on Linux.","solutions":["Cast explicitly: jax.random.random_wrap(arr.astype(jnp.uint32), impl=...)","Prefer creating keys from seeds via jax.random.key(seed)","When x64 mode is on, be aware asarray produces int64 and always cast to uint32 before wrapping"],"exampleFix":"// before\nkey = jax.random.random_wrap(jnp.array([1, 2]), impl='threefry2x32')\n\n// after\nkey = jax.random.random_wrap(jnp.array([1, 2], dtype=jnp.uint32), impl='threefry2x32')","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\ndata = jnp.asarray(data)\nif data.dtype != jnp.uint32:\n    data = data.astype(jnp.uint32)","typeGuard":"import jax.numpy as jnp\ndef is_uint32_key_data(x) -> bool:\n    return x.dtype == jnp.uint32","tryCatchPattern":null,"preventionTips":["Always specify dtype=jnp.uint32 when creating key data","In x64 mode, remember asarray defaults to int64 and cast explicitly"],"tags":["jax","prng","dtype","uint32"],"backgroundTag":"jax-prng-key-validation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}