{"record":{"id":"a9e1e212326a8128","repo":"jax-ml/jax","slug":"philox2x32-random-bits-got-invalid-prng-key","errorCode":null,"errorMessage":"philox2x32_random_bits got invalid prng key.","messagePattern":"philox2x32_random_bits got invalid prng key\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/philox2x32.py","lineNumber":188,"sourceCode":"def philox2x32_fold_in(key: typing.Array, data: typing.Array) -> typing.Array:\n  \"\"\"Fold-in an integer value to create a new Philox2x32 key.\"\"\"\n  assert not data.shape\n  return _philox2x32_fold_in(key, jnp.asarray(data, dtype=\"uint32\"))\n\n\n@api.jit\ndef _philox2x32_fold_in(key: typing.Array, data: typing.Array) -> typing.Array:\n  \"\"\"Internal implementation of philox2x32_fold_in.\"\"\"\n  out0, _ = philox2x32_p.bind(key[0], np.uint32(0), data)\n  return jnp.array([out0], dtype=np.uint32)\n\n\ndef philox2x32_random_bits(\n    key: typing.Array, bit_width: int, shape: tuple[int, ...]\n) -> typing.Array:\n  \"\"\"Sample uniform random bits using a Philox 2x32 key.\"\"\"\n  if not _is_philox2x32_key(key):\n    raise TypeError(\"philox2x32_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 _philox2x32_random_bits(key, bit_width, shape)\n\n\n@api.jit(static_argnums=(1, 2), inline=True)\ndef _philox2x32_random_bits(\n    key: typing.Array, bit_width: int, shape: tuple[int, ...]\n) -> typing.Array:\n  \"\"\"Internal implementation of philox2x32_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  counts1, counts2 = prng.iota_2x32_shape(shape)\n  out0, out1 = philox2x32_p.bind(key[0], counts1, counts2)\n\n  dtype = prng.UINT_DTYPES[bit_width]\n  if bit_width == 64:","sourceCodeStart":170,"sourceCodeEnd":206,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/philox2x32.py#L170-L206","documentation":"philox2x32_random_bits checks its key with _is_philox2x32_key (dtype uint32 and trailing shape (2,)); any other key (threefry, a plain array, or an upgraded legacy key) raises this TypeError. The function is the low-level bit-generator behind the philox2x32 implementation.","triggerScenarios":"Calling random.philox2x32_random_bits(threefry_key, 32, shape), passing a jax.random.KeyArray of a different impl, or passing a raw uint32 array whose last axis is not size 2.","commonSituations":"Mixing PRNG implementations after jax.random.use_prng_impl or upgrading old code that manipulated raw key arrays instead of typed keys.","solutions":["Create the key with the matching impl: jax.random.PRNGKey(seed, impl='philox2x32')","Convert typed keys with jax.random.key_data / raw keys via jax.random.wrap_key_data if interoperating with legacy uint32 arrays","Prefer high-level jax.random.bits API which dispatches on the key's impl"],"exampleFix":"# before\nkey = jax.random.PRNGKey(0, impl='threefry2x32')\nbits = random.philox2x32_random_bits(key, 32, (4,))\n# after\nkey = jax.random.PRNGKey(0, impl='philox2x32')\nbits = random.philox2x32_random_bits(key, 32, (4,))","handlingStrategy":"type-guard","validationCode":"key = jax.random.PRNGKey(0, impl='philox2x32')\nassert _is_philox2x32_key(key)  # or key.impl == 'philox2x32' for typed keys","typeGuard":"def is_philox2x32(k) -> bool:\n    import jax.numpy as jnp\n    d = jnp.asarray(k)\n    return d.dtype == jnp.uint32 and d.shape[-1:] == (2,)","tryCatchPattern":null,"preventionTips":["Prefer the impl-generic jax.random.bits API","Check key.impl before low-level random_bits 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"}