jax-ml/jax · error · ValueError

{mask.shape=} does not match expected shape {expected_shape}

Error message

{mask.shape=} does not match expected shape {expected_shape}

What it means

When an optional mask is given to scatter, its shape must equal the expected indexed shape (same as the value x). A mask of any other shape is rejected.

Source

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


@scatter_p.def_effectful_abstract_eval
def _scatter_abstract_eval(*flat_args, tree, add):
  ref, transforms, indices, x, mask = jax.tree.unflatten(tree, flat_args)
  if transforms:
    ref = state_types.TransformedRef(ref, transforms)
  if ref.dtype not in (jnp.int32, jnp.float32):
    raise TypeError(f"ref.dtype={ref.dtype} must be int32 or float32")
  expected_shape = _indexed_shape(ref, indices)
  if x.shape != expected_shape:
    raise ValueError(
        f"{x.shape=} does not match expected shape {expected_shape}"
    )
  if x.dtype != ref.dtype:
    raise TypeError(f"val.dtype={x.dtype} != ref.dtype={ref.dtype}")
  if mask is not None:
    if mask.shape != expected_shape:
      raise ValueError(
          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):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Reshape/broadcast the mask explicitly to expected_shape before the call
If all elements should be written, pass mask=None
Recompute the mask inside the kernel from indices if shapes differ

Example fix

// before
sc_primitives.store_scatter(ref, idx, x, mask=m)  # m.shape=(N,1)

// after
m = jnp.broadcast_to(m, x.shape)
sc_primitives.store_scatter(ref, idx, x, mask=m)

Defensive patterns

Strategy: validation

Validate before calling

if mask is not None:
    assert mask.shape == x.shape, (mask.shape, x.shape)
    mask = jnp.broadcast_to(mask, x.shape)

Type guard

def mask_shape_ok(mask, x) -> bool:
    return mask is None or tuple(mask.shape) == tuple(x.shape)

Prevention

Normalize masks with broadcast_to at creation
Pass mask=None when storing everything
Keep mask generation adjacent to the value it masks

When it happens

Trigger: Passing a mask broadcastable to (but not equal to) the selected region, e.g. mask of shape (N,1) while the region is (N,M), or a scalar/all-ones mask with wrong dims.

Common situations: Reusing a padding mask computed elsewhere with different axes kept; assuming NumPy broadcasting rules apply.

Related errors

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