jax-ml/jax · error · NotImplementedError

Get does not support masked scalar loads

Error message

Get does not support masked scalar loads

What it means

Masked scalar loads are not implemented in the SC lowering: when the load result is 0-d, passing a mask to Get raises immediately. Use an unmasked scalar load or handle bounds manually.

Source

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

    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],
      [1 if squeeze else s for s, squeeze in zip(sizes, squeeze_dims)],
  )
  memref_vec_type = ir.VectorType.get(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Clamp indices instead of masking for scalar loads
2. Load a 1-element masked vector then extract: v = masked_load(...); v[0] after ensuring valid lanes

Example fix

# before
v = ref[i, j]  # with mask applied via pallas_get(mask=...)
# after
i = jnp.clip(i, 0, n - 1)
v = ref[i, j]  # unmasked, index clamped

if scalar and mask is not None: idx = jnp.clip(idx, 0, n - 1)

null

null

When it happens

Trigger: v = ref[i, j] with mask=... (via pallas_get with a mask) where the result is a scalar, in a SparseCore kernel.

Common situations: Boundary-guarded lookups using masks, ported from TensorCore kernels.

Related errors

• masked load_p
• 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

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