jax-ml/jax · error · ValueError

Expected value and mask to have the same shape, but got valu

Error message

Expected value and mask to have the same shape, but got value shape {val_aval.shape} vs. mask shape {mask_aval.shape}.

What it means

Raised when a masked store's value shape differs from its mask shape. The TPU masked-swap lowering applies the mask lane-for-lane, so the mask must broadcast-exactly match the stored value's shape.

Source

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

@register_lowering_rule(primitives.swap_p, ensure_mlir_values=False)
def _masked_swap_lowering_rule(
    ctx: LoweringRuleContext, *args_flat, args_tree, **_
):
  ref, transforms, val, mask = args_tree.unflatten(args_flat)
  ref_aval, transforms_avals, val_aval, mask_aval = args_tree.unflatten(
      ctx.avals_in
  )
  prev_transforms, idx = _canonicalize_transforms_to_indexer(
      ref_aval, transforms, transforms_avals
  )

  if mask is not None:
    if  val_aval.dtype.itemsize != 4:
      raise NotImplementedError("masked swap with non-32-bit data")
    if val_aval.shape != mask_aval.shape:
      raise ValueError(
          "Expected value and mask to have the same shape, but got"
          f" value shape {val_aval.shape} vs. mask shape {mask_aval.shape}."
      )

  ref_block_shape, *_ = ctx.block_shapes
  ref, ref_block_shape = _transform_ref(
      ref, ref_aval, ref_block_shape, prev_transforms
  )

  ref_type = ir.MemRefType(ref.type)
  memory_space = str(ref_type.memory_space)
  is_smem_store = memory_space == "#tpu.memory_space<smem>"
  is_vmem_store = memory_space == "#tpu.memory_space<vmem>"
  (aval_out,) = ctx.avals_out
  if not isinstance(val, ir.Value):
    val = ir_constant(val, mlir_type=_dtype_to_ir_type(val_aval.dtype))
  if not is_smem_store and not ref_block_shape:
    raise NotImplementedError(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the mask the same shape as the value: mask = jnp.broadcast_to(mask, val.shape)
2. Check that the BlockSpec block shapes for value and mask refs agree
3. Recompute the mask from the current block indices (pl.program_id / block_start)

Example fix

# before
pl.store(out_ref, val, mask=mask)  # mask shape (8, 128), val shape (128, 8)
# after
pl.store(out_ref, val, mask=jnp.broadcast_to(mask.T, val.shape))

val, mask = jnp.broadcast_arrays(val, mask)
assert val.shape == mask.shape
pl.store(ref, val, mask=mask)

When it happens

Trigger: pl.store(ref, val, mask=m) with val.shape != m.shape, e.g. mask computed from a different block shape or with extra/missing dims.

Common situations: Reusing a mask computed for a different block shape; forgetting to broadcast/reshape the mask after squeezing a dim; off-by-one in BlockSpec block shapes.

Related errors

• Can only load scalars from SMEM
• masked swap with non-32-bit data
• SMEM store does not support masks
• Can only store scalars to SMEM
• masked swap with strided store

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