jax-ml/jax · error · NotImplementedError

Cannot split a Pallas key. Use fold_in instead to generate n

Error message

Cannot split a Pallas key. Use fold_in instead to generate new keys.

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic/random.py:84

  seed_data = jnp.zeros(tpu_key_impl.key_shape, dtype=jnp.int32)
  return (seed_data + seed).astype(jnp.uint32)  # Broadcast the seed.

def _random_bits(key: typing.Array, bit_width: int, shape: Shape):
  if bit_width != 32:
    raise ValueError("Bit width must be 32")
  prng_seed(key)
  return prng_random_bits(shape)

def _fold_in(key: jax_prng.PRNGKeyArray, data: typing.Array):
  key0, key1 = unwrap_pallas_seed(key)
  # Perform a cheap mixing of data into the key.
  key1 = key1 + data
  [key0, key1] = threefry2x32.apply_round([key0, key1], 13)
  return wrap_pallas_seed(key0, key1, impl="pallas_tpu")

def _split(key: typing.Array, shape: Shape):
  del key, shape
  raise NotImplementedError(
      "Cannot split a Pallas key. Use fold_in instead to generate new keys."
  )

tpu_key_impl = jax_prng.PRNGImpl(
    key_shape=(1, 2),
    seed=_seed_func,
    split=_split,
    random_bits=_random_bits,
    fold_in=_fold_in,
    name="pallas_tpu",
    tag="pl",
)
jax_prng.register_prng(tpu_key_impl)

# Implementation of the stateful Pallas PRNG API.
# Users should set the seed using the `set_seed` function,
# and call the appropriate stateful sampling functions.
# The actual key impl should never be used. The impl

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use random.fold_in(key, data) with distinct data (e.g. block index) to derive per-block keys
2. Pre-split keys on the host with the default impl and pass derived keys as kernel arguments

Example fix

# before
keys = jax.random.split(key, num_blocks)
# after
k = random.fold_in(key, block_index)

key_impl = getattr(key, '_impl', None)
if key_impl is not None and 'pallas' in str(key_impl).lower():
    # avoid split; use fold_in instead
    ...

try:
    keys = jax.random.split(key, n)
except NotImplementedError:
    keys = [jax.random.fold_in(key, i) for i in range(n)]

When it happens

Trigger: Calling jax.random.split(pallas_key) inside a Pallas TPU kernel, or any API that internally splits keys (e.g. some random distributions).

Common situations: Porting host code that pre-splits keys into per-block keys; libraries that call split implicitly on the PRNG impl.

Related errors

• PRNG keys must be loaded from SMEM. Did you set the memory s
• Bit width must be 32
• `poisson` with method='exact' is only implemented for the th
• masked load_p
• run_scoped_p with collective axes is not supported

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/b509dc0eb1c9f823. Report an issue: GitHub.