jax-ml/jax · error · ValueError

PRNG keys must be loaded from SMEM. Did you set the memory s

Error message

PRNG keys must be loaded from SMEM. Did you set the memory space to MemorySpace.SMEM in the BlockSpec for the PRNG key input?

What it means

Raised when loading PRNG keys in a TPU Pallas kernel where the key's BlockSpec memory space is not SMEM. The Pallas random implementation requires keys to reside in scalar memory (SMEM) because they are loaded one scalar at a time via memref loads. Loading keys from VMEM or ANY space is unsupported.

Source

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

  if isinstance(ref_aval.memory_space, tpu_core.AccMemorySpace):
    raise ValueError(
        "Loading from an accumulator is not supported. Use `matmul_pop` "
        "instead, which will additionally zero out the accumulator."
    )

  ref_block_shape, *_ = ctx.block_shapes
  ref, ref_block_shape = _transform_ref(
      ref, ref_aval, ref_block_shape, prev_transforms
  )
  ref_type = ir.MemRefType(ref.type)
  is_smem_load = str(ref_type.memory_space) == "#tpu.memory_space<smem>"
  (aval_out,) = ctx.avals_out
  if isinstance(aval_out.dtype, prng.KeyTy) and pl_random.is_pallas_impl(
      aval_out.dtype._impl
  ):
    # TODO(justinfu): Merge this with standard extended dtype handling.
    if not is_smem_load:
      raise ValueError("PRNG keys must be loaded from SMEM. Did you set "
                       "the memory space to MemorySpace.SMEM in the "
                       "BlockSpec for the PRNG key input?")
    return _prng_key_load_lowering_rule(ctx, *args_flat, args_tree=args_tree)
  if should_physicalize_dtype(aval_out.dtype):
    # pyrefly: ignore[bad-argument-type]
    physical_element_aval = jax_core.physical_element_aval(aval_out.dtype)
    idx = cast(NDIndexer, idx)
    if idx.int_indexer_shape:
      raise NotImplementedError()
    elt_slices = [
        indexing.Slice(0, size) for size in physical_element_aval.shape]
    idx = NDIndexer(
        indices=idx.indices + tuple(elt_slices),
        shape=idx.shape + physical_element_aval.shape,
        int_indexer_shape=(),
    )
    physical_out_dtype = physical_element_aval.dtype
    physical_out_shape = jax_core.physical_shape(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set memory_space=pltpu.MemorySpace.SMEM on the BlockSpec for the PRNG key input
2. Verify you are using the pallas-specific random implementation (jax.pallas random key type) and its required layout

Example fix

# before
keys_spec = pltpu.BlockSpec((1, 2))
# after
keys_spec = pltpu.BlockSpec((1, 2), memory_space=pltpu.MemorySpace.SMEM)

# Before launching, assert the keys BlockSpec uses SMEM
assert keys_spec.memory_space is not None and 'SMEM' in str(keys_spec.memory_space), 'PRNG keys need memory_space=SMEM'

def key_spec_is_smem(spec) -> bool:
    ms = getattr(spec, 'memory_space', None)
    return ms is not None and str(ms).endswith('SMEM')

When it happens

Trigger: Calling random_draw / pl.random with a key input whose BlockSpec (or pltpu.BlockSpec) does not set memory_space=tpu.MemorySpace.SMEM.

Common situations: Forgetting the memory_space argument in the BlockSpec for the keys argument when using pallas random RNG on TPU; defaults put blocks in VMEM.

Related errors

• Bit width must be 32
• Cannot split a Pallas key. Use fold_in instead to generate n
• Indexing into a ()-shaped Ref not yet supported on TPU.
• Can only load scalars from SMEM
• Seed key_data must be 1D.

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