{"record":{"id":"fe08704e456e440c","repo":"jax-ml/jax","slug":"rbg-random-bits-got-invalid-prng-key","errorCode":null,"errorMessage":"_rbg_random_bits got invalid prng key.","messagePattern":"_rbg_random_bits got invalid prng key\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/rbg.py","lineNumber":61,"sourceCode":"\ndef _rbg_split(key: typing.Array, shape: prng.Shape) -> typing.Array:\n  if config.threefry_partitionable.value:\n    _threefry_split = threefry2x32._threefry_split_foldlike\n  else:\n    _threefry_split = threefry2x32._threefry_split_original\n  halfkeys = key.reshape(2, 2)\n  return api.vmap(\n      _threefry_split, (0, None), len(shape))(halfkeys, shape).reshape(\n          *shape, 4)\n\ndef _rbg_fold_in(key: typing.Array, data: typing.Array) -> typing.Array:\n  assert not data.shape\n  return api.vmap(threefry2x32._threefry_fold_in, (0, None), 0)(key.reshape(2, 2), data).reshape(4)\n\ndef _rbg_random_bits(key: typing.Array, bit_width: int, shape: Sequence[int]\n                     ) -> typing.Array:\n  if not key.shape == (4,) and key.dtype == np.dtype('uint32'):\n    raise TypeError(\"_rbg_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  _, bits = lax.rng_bit_generator(key, shape, dtype=prng.UINT_DTYPES[bit_width])\n  return bits\n\nrbg_prng_impl = prng.PRNGImpl(\n    key_shape=(4,),\n    seed=_rbg_seed,\n    split=_rbg_split,\n    random_bits=_rbg_random_bits,\n    fold_in=_rbg_fold_in,\n    name='rbg',\n    tag='rbg')\n\nprng.register_prng(rbg_prng_impl)\n\n\ndef _unsafe_rbg_split(key: typing.Array, shape: prng.Shape) -> typing.Array:","sourceCodeStart":43,"sourceCodeEnd":79,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/rbg.py#L43-L79","documentation":"The RBG (stateless random-bit generator) implementation stores each key as exactly 4 uint32 words. This TypeError from _rbg_random_bits fires when the key argument does not have shape (4,) and dtype uint32 — typically because a threefry-style (2,)-word key or an ordinary array was passed to an rbg-based operation.","triggerScenarios":"Calling jax.random.bits or sampling functions with impl='rbg' on a threefry key; passing lax.rng_bit_generator a key of the wrong family; hand-building rbg key data with 2 words instead of 4.","commonSituations":"Mixing keys created with jax.random.key(seed) (threefry default) into code configured for rbg; configuration changes of the global PRNG implementation without regenerating stored keys.","solutions":["Create the key with matching impl: jax.random.key(seed, impl='rbg')","Validate key shape/dtype before use: key.shape == (4,) and key.dtype == jnp.uint32 (via key_data for typed keys)","Regenerate any persisted keys after changing implementations"],"exampleFix":"// before\nkey = jax.random.key(0)  # threefry, shape (2,)\nbits = jax.random.bits(key, shape=(8,), dtype=jnp.uint32, impl='rbg')\n\n// after\nkey = jax.random.key(0, impl='rbg')\nbits = jax.random.bits(key, shape=(8,), dtype=jnp.uint32, impl='rbg')","handlingStrategy":"validation","validationCode":"import jax, jax.numpy as jnp\ndata = jax.random.key_data(key)\nassert data.shape == (4,) and data.dtype == jnp.uint32, f'bad rbg key: {data.shape} {data.dtype}'","typeGuard":"import jax, jax.numpy as jnp\ndef is_rbg_key(key) -> bool:\n    d = jax.random.key_data(key)\n    return d.shape == (4,) and d.dtype == jnp.uint32","tryCatchPattern":null,"preventionTips":["Create keys with jax.random.key(seed, impl='rbg') for rbg code paths","Never reuse threefry keys with rbg-based samplers"],"tags":["jax","prng","rbg","shape-mismatch"],"backgroundTag":"jax-prng-key-validation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}