{"record":{"id":"60d35695ae8b927d","repo":"jax-ml/jax","slug":"only-32-bit-prng-supported","errorCode":null,"errorMessage":"Only 32-bit PRNG supported.","messagePattern":"Only 32-bit PRNG supported\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/random/philox.py","lineNumber":193,"sourceCode":"        key, padded_shape, shape,\n        block_size=BLOCK_SIZE, offset=offset,\n        fuse_output=fuse_output)\n    return padded_result[..., :shape[-2], :shape[-1]]\n  else:\n    return philox_4x32_kernel(key, shape, shape,\n                              block_size=BLOCK_SIZE, offset=offset,\n                              fuse_output=fuse_output)\n\n\ndef philox_split(key, shape: Shape):\n  \"\"\"Splits the key into two keys of the same shape.\"\"\"\n  bits1, bits2 = philox_4x32_count(key, shape, fuse_output=False)\n  return jnp.stack([bits1, bits2], axis=bits1.ndim)\n\n\ndef philox_random_bits(key, bit_width: int, shape: Shape):\n  if bit_width != 32:\n    raise ValueError(\"Only 32-bit PRNG supported.\")\n  return philox_4x32_count(key, shape, fuse_output=True)\n\n\ndef philox_fold_in(key, data):\n  assert data.ndim == 0\n  return philox_4x32_count(key, (), offset=data, fuse_output=False)\n\n\nplphilox_prng_impl = prng.PRNGImpl(\n    key_shape=(2,),\n    seed=threefry2x32.threefry_seed,\n    split=philox_split,\n    random_bits=philox_random_bits,\n    fold_in=philox_fold_in,\n    name=\"pallas_philox4x32\",\n    tag=\"pllox\")\n\nprng.register_prng(plphilox_prng_impl)","sourceCodeStart":175,"sourceCodeEnd":211,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/random/philox.py#L175-L211","documentation":"The TPU Pallas Philox implementation only generates 32-bit random values: philox_random_bits raises ValueError for any bit_width other than 32 because the kernel produces uint32 words only.","triggerScenarios":"Calling philox_random_bits(key, 8, shape) or (key, 16, shape) or (key, 64, shape).","commonSituations":"Plugging this backend into a generic PRNG interface (e.g. jax.random.bits with dtype uint8/uint16/uint64) where bit width varies by dtype.","solutions":["Request 32 bits and truncate/shift yourself for smaller widths: (bits >> 24).astype(jnp.uint8)","For 64-bit, generate two 32-bit halves and combine: (hi.astype(uint64) << 32) | lo","Route non-32-bit requests to the standard jax.random implementation instead"],"exampleFix":"// before\nu8 = philox_random_bits(key, 8, shape)\n// after\nu8 = (philox_random_bits(key, 32, shape) >> 24).astype(jnp.uint8)","handlingStrategy":"validation","validationCode":"assert bit_width == 32, 'TPU pallas philox only supports 32-bit; derive smaller widths by shifting'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive uint8/uint16 by shifting 32-bit output","Dispatch other bit widths to jax.random's default backend"],"tags":["jax","pallas","tpu","random","philox","bit-width"],"backgroundTag":"unsupported-bit-width","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}