jax-ml/jax · error · NotImplementedError

Can only load a single key per load.

Error message

Can only load a single key per load.

What it means

Raised in _prng_key_load_lowering_rule: each Pallas load may fetch at most one PRNG key, so a key ref/block whose leading dimension permits multiple keys per load is rejected with NotImplementedError; split the load into one key at a time.

Solutions

1. Load one key at a time (scalar indexer) and loop or split keys inside the kernel
2. Use random_split/random_fold_in to derive multiple keys from the single loaded seed

# in-kernel: load keys[0:1, 0:1]-style scalar indexers only
assert all(s == 1 for s in idx_shape), 'load one key at a time'

When it happens

Trigger: pl.load on an SMEM PRNG key ref with a slice/index whose shape contains dims > 1, requesting multiple keys in one load.

Common situations: Slicing the key block with a range (e.g. keys[0:4]) hoping to get several keys at once; vectorizing RNG inside the kernel.

Related errors

• Bit width must be 32
• Cannot do int indexing on TPU
• Cannot split a Pallas key. Use fold_in instead to generate…
• Dimension must be 0 for 1D iota.
• Integer indexing of refs that follows a non-trivial slice…

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

Appendix: source

Thrown at jax/_src/pallas/mosaic/lowering.py:2402

  )
  prev_transforms, idx = _canonicalize_transforms_to_indexer(
      ref_aval, transforms, transforms_avals
  )
  (aval_out,) = ctx.avals_out
  assert isinstance(aval_out.dtype, prng.KeyTy)
  key_shape = aval_out.dtype._impl.key_shape
  ref_block_shape, *_ = ctx.block_shapes
  idx = cast(NDIndexer, idx)
  ref, ref_block_shape = _transform_ref(
      ref, ref_aval, ref_block_shape, prev_transforms
  )

  if len(key_shape) != 2:
    raise NotImplementedError("Seed key_data must be 1D.")
  if key_shape[0] != 1:
    raise NotImplementedError("Leading dimension of seed key_data must be 1.")
  if not all(s == 1 for s in idx.shape):
    raise NotImplementedError("Can only load a single key per load.")
  assert ref_block_shape[-2:] == key_shape, f"{ref_block_shape=} {key_shape=}"

  load_ops = []
  for i in range(key_shape[1]):
    ref_shape = tuple(
        dim for dim in ref_block_shape if dim is not pallas_core.squeezed
    )
    scalar_idx = NDIndexer(
        indices=(*idx.indices, 0, i), shape=ref_shape, int_indexer_shape=()
    )
    starts, _, _, _, _ = _indexer_to_start_size_stride(
        scalar_idx,
        ref_block_shape,
        cast_to_index=True,
    )
    load_ops.append(memref.load(ref, starts))
  return KeyScalarBundle(scalars=load_ops, key_shape=tuple(key_shape))

View on GitHub (pinned to 1e1c6a8fc0)