{"record":{"id":"00b6c29c7dae8b54","repo":"jax-ml/jax","slug":"threefry-random-bits-got-invalid-prng-key","errorCode":null,"errorMessage":"threefry_random_bits got invalid prng key.","messagePattern":"threefry_random_bits got invalid prng key\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/threefry2x32.py","lineNumber":324,"sourceCode":"  k1, k2 = key\n  counts1, counts2 = prng.iota_2x32_shape(shape)\n  bits1, bits2 = threefry2x32_p.bind(k1, k2, counts1, counts2)\n  return jnp.stack([bits1, bits2], axis=bits1.ndim)\n\n\ndef threefry_fold_in(key: typing.Array, data: typing.Array) -> typing.Array:\n  assert not data.shape\n  return _threefry_fold_in(key, jnp.asarray(data, dtype='uint32'))\n\n@api.jit\ndef _threefry_fold_in(key, data):\n  return threefry_2x32(key, threefry_seed(data))\n\n\ndef threefry_random_bits(key: typing.Array, bit_width, shape):\n  \"\"\"Sample uniform random bits of given width and shape using PRNG key.\"\"\"\n  if not _is_threefry_prng_key(key):\n    raise TypeError(\"threefry_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\n  if config.threefry_partitionable.value:\n    return _threefry_random_bits_partitionable(key, bit_width, shape)\n  else:\n    return _threefry_random_bits_original(key, bit_width, shape)\n\ndef _threefry_random_bits_partitionable(key: typing.Array, bit_width, shape):\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  k1, k2 = key\n  counts1, counts2 = prng.iota_2x32_shape(shape)\n  bits1, bits2 = threefry2x32_p.bind(k1, k2, counts1, counts2)\n\n  dtype = prng.UINT_DTYPES[bit_width]\n  if bit_width == 64:","sourceCodeStart":306,"sourceCodeEnd":342,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/threefry2x32.py#L306-L342","documentation":"threefry_random_bits requires its key argument to be a valid threefry PRNG key: raw uint32 data of shape (2,) (or batched equivalents). This TypeError fires when the key is from another implementation (e.g. rbg with 4 words), already-consumed raw data of the wrong shape, or a non-key array.","triggerScenarios":"Passing an rbg key to threefry-based threefry_random_bits (reachable via jax.random.bits with mismatched impls); calling the internal API with arbitrary uint32 arrays; mixing keys across generator families after config changes.","commonSituations":"Global PRNG implementation switched to/from rbg without regenerating keys; code paths that conditionally use _double_threefry_random_bits with keys of the wrong family; custom samplers calling internal bit-generators directly.","solutions":["Generate keys with the matching implementation: jax.random.key(seed, impl='threefry2x32')","Validate raw key data shape (2,) and uint32 dtype before calling","Use public jax.random.bits/rand APIs, which dispatch on the key's impl"],"exampleFix":"// before\nkey = jax.random.key(0, impl='rbg')  # 4-word key\nbits = jax.random.bits(key, shape=(8,), dtype=jnp.uint32)  # routed to threefry path\n\n// after\nkey = jax.random.key(0)  # threefry2x32\nbits = jax.random.bits(key, shape=(8,), dtype=jnp.uint32)","handlingStrategy":"type-guard","validationCode":"import jax, jax.numpy as jnp\ndata = jax.random.key_data(key)\nassert data.shape[-2:] == (2,) and data.dtype == jnp.uint32, 'not a threefry key'","typeGuard":"import jax, jax.numpy as jnp\ndef is_threefry_key(key) -> bool:\n    d = jax.random.key_data(key)\n    return d.dtype == jnp.uint32 and d.shape[-2:] == (2,)","tryCatchPattern":null,"preventionTips":["Create keys with the default impl (threefry2x32) for threefry code paths","Use jax.random.bits/rand public APIs which dispatch on key impl"],"tags":["jax","prng","threefry","key-validation"],"backgroundTag":"jax-prng-key-validation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}