jax-ml/jax · error · NotImplementedError

Leading dimension of seed key_data must be 1.

Error message

Leading dimension of seed key_data must be 1.

What it means

Raised when loading PRNG seed key_data whose leading dimension is not 1. The TPU lowering expects exactly one key per load with layout (1, num_elems); a batched leading dimension is unimplemented.

Source

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

  ref_aval, transforms_avals, _, _ = args_tree.unflatten(
      ctx.avals_in
  )
  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))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Load a single key of shape (1, 2) and derive per-use keys in-kernel (split/fold_in)

# host side: ensure the key argument is a single key, not a batch
assert key.ndim == 0 or key.shape == (1,), 'load one key of shape (1, 2), not a batch'

When it happens

Trigger: pl.load of a PRNG key ref where key_shape[0] != 1, e.g. loading a (k, 2) array of keys as the seed.

Common situations: Trying to load multiple keys at once for per-lane RNG instead of splitting keys inside the kernel with random_fold_in/split.

Related errors

• Seed key_data must be 1D.
• PRNG keys must be loaded from SMEM. Did you set the memory s
• Can only load a single key per load.
• Bit width must be 32
• Cannot split a Pallas key. Use fold_in instead to generate n

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