{"record":{"id":"f41ad465cce4da74","repo":"jax-ml/jax","slug":"bit-width-must-be-32","errorCode":null,"errorMessage":"Bit width must be 32","messagePattern":"Bit width must be 32","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/random.py","lineNumber":71,"sourceCode":"\n  if vmapped_key:\n    pallas_key_data = jax.vmap(generate_key)(key)\n  else:\n    pallas_key_data = generate_key(key)\n  return jax_api_random.wrap_key_data(pallas_key_data, impl=\"pallas_tpu\")\n\ndef is_pallas_impl(impl: jax_prng.PRNGImpl) -> bool:\n  \"\"\"Returns True if the PRNGImpl is a Pallas-specific implementation.\"\"\"\n  return impl == tpu_key_impl or impl == tpu_internal_stateful_impl\n\n\ndef _seed_func(seed: jnp.int32):\n  seed_data = jnp.zeros(tpu_key_impl.key_shape, dtype=jnp.int32)\n  return (seed_data + seed).astype(jnp.uint32)  # Broadcast the seed.\n\ndef _random_bits(key: typing.Array, bit_width: int, shape: Shape):\n  if bit_width != 32:\n    raise ValueError(\"Bit width must be 32\")\n  prng_seed(key)\n  return prng_random_bits(shape)\n\ndef _fold_in(key: jax_prng.PRNGKeyArray, data: typing.Array):\n  key0, key1 = unwrap_pallas_seed(key)\n  # Perform a cheap mixing of data into the key.\n  key1 = key1 + data\n  [key0, key1] = threefry2x32.apply_round([key0, key1], 13)\n  return wrap_pallas_seed(key0, key1, impl=\"pallas_tpu\")\n\ndef _split(key: typing.Array, shape: Shape):\n  del key, shape\n  raise NotImplementedError(\n      \"Cannot split a Pallas key. Use fold_in instead to generate new keys.\"\n  )\n\ntpu_key_impl = jax_prng.PRNGImpl(\n    key_shape=(1, 2),","sourceCodeStart":53,"sourceCodeEnd":89,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/random.py#L53-L89","documentation":"The Pallas TPU random-bit generator is hardwired to 32-bit output, so _random_bits raises ValueError when bit_width != 32. This is called through the PRNGImpl when random_bits/normal/etc. request a different width inside a Pallas kernel.","triggerScenarios":"Using jax.random bits APIs inside a Pallas TPU kernel with the pallas_tpu PRNG impl where the requested bit width is not 32 (e.g. random bits of width 8 or 64).","commonSituations":"Calling random-level APIs that default to the dtype width (uint16/uint64 keys) inside a kernel; porting host-side jax.random code into Pallas.","solutions":["Request 32-bit random values inside Pallas kernels (uint32/float32 samples)","Generate non-32-bit values on the host or derive them from 32-bit samples"],"exampleFix":"# before\nbits = random_bits(key, bit_width=8, shape=shape)\n# after\nbits = (random_bits(key, bit_width=32, shape=shape) & 0xFF).astype(jnp.uint8)","handlingStrategy":"validation","validationCode":"assert bit_width == 32, 'Pallas TPU PRNG only supports 32-bit'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Sample uint32/f32 inside kernels; derive narrower dtypes via masking/casting"],"tags":["jax","pallas","tpu","random","prng"],"backgroundTag":"unsupported-argument-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}