jax-ml/jax · error · NotImplementedError

Indexing into a ()-shaped Ref not yet supported on TPU.

Error message

Indexing into a ()-shaped Ref not yet supported on TPU.

What it means

Raised when a Pallas TPU load targets a ()-shaped (scalar) Ref that is not in SMEM. VMEM ops on TPU are vector ops and cannot address a scalar block, so scalar loads are only supported from SMEM; loading a scalar VMEM block is not yet implemented.

Source

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

    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(
        aval_out.shape, aval_out.dtype
    )
  else:
    physical_out_dtype = aval_out.dtype
    physical_out_shape = aval_out.shape
  if not is_smem_load and not ref_block_shape:
    raise NotImplementedError(
        "Indexing into a ()-shaped Ref not yet supported on TPU.")
  starts, sizes, strides, _, _ = _indexer_to_start_size_stride(
      idx,
      ref_block_shape,
      cast_to_index=True,
  )
  need_stride = not all((s is None or s == 1) for s in strides)
  if is_smem_load:
    if ctx.avals_out[0].shape:
      raise ValueError("Can only load scalars from SMEM")
    return _maybe_cast_load_to_bool(ctx, aval_out, memref.load(ref, starts))
  elif str(ref_type.memory_space) != "#tpu.memory_space<vmem>":
    extra = ""
    if str(ref_type.memory_space) == "#tpu.memory_space<any>":
      extra = " ANY memory space can only be accessed using async_copy."
    raise ValueError(
        "Loads are only allowed on VMEM and SMEM references." + extra
    )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use memory_space=pltpu.MemorySpace.SMEM for the scalar input's BlockSpec
Or give the value shape (1,) instead of () and keep it in VMEM
Pass scalars as compile-time constants or via inlined constants rather than refs

Example fix

# before
spec = pltpu.BlockSpec(())
# after
spec = pltpu.BlockSpec((), memory_space=pltpu.MemorySpace.SMEM)

Defensive patterns

Strategy: validation

Validate before calling

def scalar_spec_ok(spec):
    block = getattr(spec, 'block_shape', None)
    ms = str(getattr(spec, 'memory_space', '') or '')
    if block == () and 'SMEM' not in ms:
        raise ValueError('scalar block must use memory_space=SMEM on TPU')

Prevention

Prefer (1,) block shapes over () for VMEM values
Route all scalar inputs through SMEM specs via a shared helper

When it happens

Trigger: pl.load on a Ref whose block shape is () (scalar block) while the ref lives in VMEM (not SMEM). Typically from a BlockSpec with scalar block shape for a scalar input.

Common situations: Passing scalars (like a loop counter or scalar hyperparameter) into the kernel with default VMEM BlockSpec; squeezing a dimension away so the block becomes scalar.

Related errors

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