{"record":{"id":"b509dc0eb1c9f823","repo":"jax-ml/jax","slug":"cannot-split-a-pallas-key-use-fold-in-instead-to","errorCode":null,"errorMessage":"Cannot split a Pallas key. Use fold_in instead to generate new keys.","messagePattern":"Cannot split a Pallas key\\. Use fold_in instead to generate new keys\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/random.py","lineNumber":84,"sourceCode":"  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),\n    seed=_seed_func,\n    split=_split,\n    random_bits=_random_bits,\n    fold_in=_fold_in,\n    name=\"pallas_tpu\",\n    tag=\"pl\",\n)\njax_prng.register_prng(tpu_key_impl)\n\n# Implementation of the stateful Pallas PRNG API.\n# Users should set the seed using the `set_seed` function,\n# and call the appropriate stateful sampling functions.\n# The actual key impl should never be used. The impl","sourceCodeStart":66,"sourceCodeEnd":102,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/random.py#L66-L102","documentation":"The Pallas TPU PRNG impl does not implement split; keys can only be derived by folding data into an existing key. Calling jax.random.split on a pallas_tpu key raises NotImplementedError with a hint to use fold_in.","triggerScenarios":"Calling jax.random.split(pallas_key) inside a Pallas TPU kernel, or any API that internally splits keys (e.g. some random distributions).","commonSituations":"Porting host code that pre-splits keys into per-block keys; libraries that call split implicitly on the PRNG impl.","solutions":["Use random.fold_in(key, data) with distinct data (e.g. block index) to derive per-block keys","Pre-split keys on the host with the default impl and pass derived keys as kernel arguments"],"exampleFix":"# before\nkeys = jax.random.split(key, num_blocks)\n# after\nk = random.fold_in(key, block_index)","handlingStrategy":"fallback","validationCode":"key_impl = getattr(key, '_impl', None)\nif key_impl is not None and 'pallas' in str(key_impl).lower():\n    # avoid split; use fold_in instead\n    ...","typeGuard":null,"tryCatchPattern":"try:\n    keys = jax.random.split(key, n)\nexcept NotImplementedError:\n    keys = [jax.random.fold_in(key, i) for i in range(n)]","preventionTips":["Prefer fold_in(key, block_index) for per-block keys in Pallas","Split keys on the host with the default impl before passing into kernels"],"tags":["jax","pallas","tpu","random","prng","not-implemented"],"backgroundTag":"operation-not-supported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}