jax-ml/jax · error · NotImplementedError

Seed key_data must be 1D.

Error message

Seed key_data must be 1D.

What it means

Raised when the key_data of a PRNG seed loaded in a TPU Pallas kernel is not effectively 1D. The lowering requires key_data shape (1, N) after accounting for the leading impl dimension; anything else (e.g. a 2D key array) is unimplemented.

Source

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

  """
  ref, transforms, _, _ = args_tree.unflatten(args_flat)
  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,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a single seed key with 1D key_data: use jax.random.key / pallas random seeding so key_data is (1, 2) or (1, n)
2. Generate multiple keys inside the kernel via random_draw from one seed instead of loading a key array

import jax
key = jax.random.key(seed)
assert key.shape == ()  # single key; key_data stays 1D (1, n)

When it happens

Trigger: pl.load on an SMEM PRNG key ref whose key_shape (from aval_out.shape plus impl layout) has len != 2, i.e. key data that isn't a flat 1D array of shape (1, n).

Common situations: Passing a shaped/random key array (e.g. shape (4, 2)) as the seed input rather than a single unraveled key; using a non-pallas RNG impl shape.

Related errors

• Leading dimension of seed key_data must be 1.
• 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/6872498fbb360190. Report an issue: GitHub.