jax-ml/jax · error · NotImplementedError

Get only supports loading scalars from SMEM.

Error message

Get only supports loading scalars from SMEM.

What it means

On SC, scalar (0-d) loads are only supported from SMEM. Trying to load a scalar from VMEM (or another space) fails; instead load a 1-element array and extract the element afterwards.

Source

Thrown at jax/_src/pallas/mosaic/sc_lowering.py:135

        " supported on SC"
    )
  if not all(s == 1 for s in strides):
    raise NotImplementedError(
        "Get only supports slices with stride 1, got {strides}"
    )

  if (out_aval.ndim == 0) != (ref_memory_space is MemorySpace.SMEM):
    message = "Get only supports loading scalars from SMEM."
    if ref_memory_space is MemorySpace.SMEM:
      message += " Trying to load an array of shape {out_aval.shape}."
    elif ref_memory_space is MemorySpace.VMEM:
      message += (
          " To load a scalar from VMEM, load an array first and then extract a"
          " particular element, e.g. ``v = ref[pl.ds(idx, ...)]; v[0]``."
      )
    else:
      message += f" Trying to load a scalar from {ref_memory_space!r}."
    raise NotImplementedError(message)
  if out_aval.ndim == 0:
    if mask is not None:
      raise NotImplementedError("Get does not support masked scalar loads")
    return memref.load(ref, starts)

  if not ctx.lowering_context.needs_layout_passes:
    _check_aval_is_supported("Get", out_aval)
  out_vec_type = ir.VectorType.get(
      out_aval.shape, _dtype_to_ir_type(out_aval.dtype)
  )
  if not ctx.lowering_context.needs_layout_passes:
    return tpu.vector_load(
        out_vec_type, ref, indices=starts, strides=[], mask=mask
    )
  # Load at the full memref rank, keeping integer-indexed dims as size 1,
  # because apply-vector-layout requires the vector rank to match the memref.
  memref_vec_shape = cast(
      Sequence[int],

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Load a length-1 dynamic slice and index element 0: v = ref[pl.ds(idx, 1)]; v[0]
2. Move the scalar into SMEM and read from there

Example fix

# before
s = vmem_ref[i, j]
# after
v = vmem_ref[pl.ds(i, 1), pl.ds(j, 1)]
s = v[0]

if out_aval.ndim == 0 and space is not MemorySpace.SMEM:
    v = ref[pl.ds(idx, 1)]  # load vector instead

null

null

When it happens

Trigger: out_aval.ndim == 0 while ref_memory_space is not SMEM, e.g. v = vmem_ref[i, j] with scalar result inside an SC kernel.

Common situations: Scalar coefficient lookups from VMEM buffers; index math that collapses all dims.

Related errors

• The current TPU does not have SparseCores
• Mesh has {self.num_cores} cores, but the current TPU chip ha
• You can't use two different ScalarSubcoreMeshes.
• {self} should have the same core axis name and number of cor
• {self} should have a different axis name from the TensorCore

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