{"record":{"id":"6d2c11c7df7e0a6e","repo":"jax-ml/jax","slug":"only-32-bit-prng-supported-6d2c11","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/threefry.py","lineNumber":92,"sourceCode":"        k1, k2, counts_hi, counts_lo)\n    out_bits = o1 ^ o2\n    out_ref[...] = out_bits.reshape(out_ref.shape)\n\n  key = key.reshape((1, 2))\n  out = jax.ShapeDtypeStruct(shape, dtype=jnp.uint32)\n  block_shape = (1,) * (len(shape)-2) + block_size\n  result = pl.pallas_call(\n      kernel,\n      in_specs=[pl.BlockSpec(memory_space=pltpu.SMEM)],\n      out_specs=pl.BlockSpec(block_shape, lambda *idxs: idxs),\n      grid=grid_dims,\n      out_shape=out,\n  )(key)\n  return result\n\ndef plthreefry_random_bits(key, bit_width: int, shape: Shape):\n  if bit_width != 32:\n    raise ValueError(\"Only 32-bit PRNG supported.\")\n  if len(shape) == 0:\n    return plthreefry_random_bits(key, bit_width, (1, 1))[0, 0]\n  elif len(shape) == 1:\n    return plthreefry_random_bits(key, bit_width, (1, *shape))[0]\n\n  requires_pad = (\n      shape[-2] % BLOCK_SIZE[-2] != 0) or (shape[-1] % BLOCK_SIZE[-1] != 0)\n  if requires_pad:\n    padded_shape = tuple(shape[:-2]) + (\n        prng_utils.round_up(shape[-2], BLOCK_SIZE[-2]),\n        prng_utils.round_up(shape[-1], BLOCK_SIZE[-1]),\n    )\n    padded_result = threefry_2x32_count(\n        key, padded_shape, shape, block_size=BLOCK_SIZE)\n    return padded_result[..., :shape[-2], :shape[-1]]\n  else:\n    return threefry_2x32_count(key, shape, shape, block_size=BLOCK_SIZE)\n","sourceCodeStart":74,"sourceCodeEnd":110,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/random/threefry.py#L74-L110","documentation":"The Pallas TPU Threefry backend only supports 32-bit random output; plthreefry_random_bits raises ValueError for bit_width != 32 because the 2x32 kernel emits uint32 words only.","triggerScenarios":"Calling plthreefry_random_bits(key, 16, shape) or any width other than 32.","commonSituations":"A generic PRNG-dispatch layer (like jax.random.bits' internal backend selection) forwarding dtypes such as uint8/uint16/uint64 to this TPU backend.","solutions":["Generate 32 bits and downcast: (bits >> 24).astype(jnp.uint8)","For 64-bit, combine two 32-bit draws","Dispatch non-32-bit widths to the default threefry PRNG implementation"],"exampleFix":"// before\nu16 = plthreefry_random_bits(key, 16, shape)\n// after\nu16 = (plthreefry_random_bits(key, 32, shape) >> 16).astype(jnp.uint16)","handlingStrategy":"validation","validationCode":"assert bit_width == 32, 'pallas threefry only emits uint32; downcast manually for smaller widths'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Downcast from 32-bit output for uint8/uint16 needs","Route non-32-bit requests to the standard threefry backend"],"tags":["jax","pallas","tpu","random","threefry","bit-width"],"backgroundTag":"unsupported-bit-width","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}