{"record":{"id":"e6f19a62d50654f2","repo":"jax-ml/jax","slug":"philox4x32-random-bits-got-invalid-prng-key","errorCode":null,"errorMessage":"philox4x32_random_bits got invalid prng key.","messagePattern":"philox4x32_random_bits got invalid prng key\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/philox4x32.py","lineNumber":216,"sourceCode":"\n\n@api.jit\ndef _philox4x32_fold_in(key: typing.Array, data: typing.Array) -> typing.Array:\n  \"\"\"Internal implementation of philox4x32_fold_in.\"\"\"\n  # Hash the key with the data used as part of the counter.\n  k0, k1 = key[0], key[1]\n  out0, out1, _, _ = philox4x32_p.bind(\n      k0, k1, np.uint32(0), np.uint32(0), np.uint32(0), data\n  )\n  return jnp.array([out0, out1], dtype=np.uint32)\n\n\ndef philox4x32_random_bits(\n    key: typing.Array, bit_width: int, shape: tuple[int, ...]\n) -> typing.Array:\n  \"\"\"Sample uniform random bits using a Philox 4x32 key.\"\"\"\n  if not _is_philox4x32_key(key):\n    raise TypeError(\"philox4x32_random_bits got invalid prng key.\")\n  if bit_width not in (8, 16, 32, 64):\n    raise TypeError(\"requires 8-, 16-, 32- or 64-bit field width.\")\n  return _philox4x32_random_bits(key, bit_width, shape)\n\n\n@api.jit(static_argnums=(1, 2), inline=True)\ndef _philox4x32_random_bits(\n    key: typing.Array, bit_width: int, shape: tuple[int, ...]\n) -> typing.Array:\n  \"\"\"Internal implementation of philox4x32_random_bits.\"\"\"\n  if all(core.is_constant_dim(d) for d in shape) and math.prod(shape) > 2**64:\n    raise NotImplementedError(\"random bits array of size exceeding 2 ** 64\")\n\n  k0, k1 = key[0], key[1]\n  counts1, counts2 = prng.iota_2x32_shape(shape)\n  zeros = jnp.zeros(shape, dtype=np.uint32)\n\n  out0, out1, out2, out3 = philox4x32_p.bind(","sourceCodeStart":198,"sourceCodeEnd":234,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/philox4x32.py#L198-L234","documentation":"philox4x32_random_bits validates that its key is a Philox 4x32 key (uint32 dtype with trailing shape (4,)); keys of any other PRNG implementation or raw arrays raise this TypeError.","triggerScenarios":"Passing a threefry or typed jax.random.KeyArray of a different impl, or a raw uint32 array with last dim != 4.","commonSituations":"Interoperability code that passes whichever key it has; upgrading legacy raw-key manipulation to typed keys.","solutions":["Create the key with impl='philox4x32' or use jax.random.wrap_key_data on correctly-shaped uint32 data","Use impl-generic jax.random.bits instead","Check the key's impl with key.impl before calling the low-level API"],"exampleFix":"# before\nbits = random.philox4x32_random_bits(threefry_key, 32, (8,))\n# after\nkey = jax.random.PRNGKey(0, impl='philox4x32')\nbits = random.philox4x32_random_bits(key, 32, (8,))","handlingStrategy":"type-guard","validationCode":"key = jax.random.PRNGKey(0, impl='philox4x32')\nassert jnp.asarray(key).shape[-1:] == (4,) and jnp.asarray(key).dtype == jnp.uint32","typeGuard":"def is_philox4x32(k) -> bool:\n    import jax.numpy as jnp\n    d = jnp.asarray(k)\n    return d.dtype == jnp.uint32 and d.shape[-1:] == (4,)","tryCatchPattern":null,"preventionTips":["Use impl-generic jax.random.bits where possible","Verify key provenance before low-level calls"],"tags":["jax","prng","philox","key-validation"],"backgroundTag":"prng-key-implementation-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}