jax-ml/jax · error · ValueError

Scatter only supports VectorSubcoreMesh, got {type(ref_aval.

Error message

Scatter only supports VectorSubcoreMesh, got {type(ref_aval.memory_space.mesh)}

What it means

When the scatter ref lives in a CoreMemorySpace, its mesh must be a VectorSubcoreMesh; other mesh types are unsupported by the SC scatter lowering.

Source

Thrown at jax/_src/pallas/mosaic/sc_primitives.py:354

          f"{mask.shape=} does not match expected shape {expected_shape}"
      )
    if mask.dtype != jnp.bool:
      raise TypeError(f"Mask must be a boolean array, got {mask.dtype}")
  effects: set[jax_core.Effect] = {state_types.WriteEffect(0)}
  if add:
    effects.add(state_types.ReadEffect(0))
  return (), effects


@sc_lowering.register_lowering_rule(scatter_p)
def _scatter_lowering_rule(
    ctx: sc_lowering.LoweringRuleContext, *flat_args, tree, add
):
  ref, transforms, indices, x, mask = jax.tree.unflatten(tree, flat_args)
  ref_aval, *_ = tree.unflatten(ctx.avals_in)
  if isinstance(ref_aval.memory_space, pallas_core.CoreMemorySpace):
    if not isinstance(ref_aval.memory_space.mesh, sc_core.VectorSubcoreMesh):
      raise ValueError(
          "Scatter only supports VectorSubcoreMesh, got"
          f" {type(ref_aval.memory_space.mesh)}"
      )
    memory_space = ref_aval.memory_space.memory_space
  else:
    memory_space = ref_aval.memory_space
  if memory_space not in (
      tpu_core.MemorySpace.VMEM,
      pallas_core.MemorySpace.DEFAULT,
  ):
    raise ValueError(
        f"Scatter only supports storing to VMEM, got {memory_space}"
    )
  if transforms:
    ref_block_shape, *_ = ctx.block_shapes
    ref, _ = tc_lowering._transform_ref(
        ref, ref_aval, ref_block_shape, transforms
    )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Construct the memory scope with sc_core.VectorSubcoreMesh (the only supported mesh)
2. Check the mesh type you pass to CoreMemorySpace in your kernel's memory declarations
3. Update/align jax version so the VectorSubcoreMesh API matches your code

Example fix

// before
ms = pallas_core.CoreMemorySpace(memory_space=..., mesh=custom_mesh)

// after
from jax._src.pallas.mosaic import sc_core
ms = pallas_core.CoreMemorySpace(memory_space=..., mesh=sc_core.VectorSubcoreMesh())

ms = ref_aval.memory_space
if isinstance(ms, pallas_core.CoreMemorySpace):
    assert isinstance(ms.mesh, sc_core.VectorSubcoreMesh), type(ms.mesh)

def has_supported_mesh(ref_aval) -> bool:
    ms = ref_aval.memory_space
    return not isinstance(ms, pallas_core.CoreMemorySpace) or isinstance(ms.mesh, sc_core.VectorSubcoreMesh)

When it happens

Trigger: Binding scatter to a ref whose memory_space is a pallas_core.CoreMemorySpace constructed with a custom/unsupported mesh (e.g. a TensorCore mesh or new mesh type).

Common situations: Experimenting with new mesh APIs in Pallas; version skew where a mesh class was renamed/split; using TC-oriented memory scopes in an SC kernel.

Related errors

• {x.shape=} does not match expected shape {expected_shape}
• val.dtype={x.dtype} != ref.dtype={ref.dtype}
• {mask.shape=} does not match expected shape {expected_shape}
• Scatter only supports storing to VMEM, got {memory_space}
• Indices must not be empty

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