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 when the index used to load PRNG keys selects more than one element, i.e. the indexer shape is not all 1s. The PRNG key load path is scalar-only: each load fetches exactly one key.

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)

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

• Cannot do int indexing on TPU
• PRNG keys must be loaded from SMEM. Did you set the memory s
• Seed key_data must be 1D.
• Leading dimension of seed key_data must be 1.
• Dimension must be 0 for 1D iota.

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