jax-ml/jax · error · ValueError

Unexpected mask shape, got: {mask.shape}, expected: {shape}

Error message

Unexpected mask shape, got: {mask.shape}, expected: {shape}

What it means

All masks inside a MultiHeadMask must share the same shape; per-head masks with different (q, kv) shapes cannot form a uniform multi-head mask, so __post_init__ rejects the first mismatch found.

Source

Thrown at jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask.py:188

  def __hash__(self):
    return hash((type(self),) + (self.left, self.right))


@dataclasses.dataclass
class MultiHeadMask(Mask):
  """Lazy multihead mask, combines multiple lazy masks one per head."""

  masks: Sequence[Mask]

  def __post_init__(self):
    if not self.masks:
      raise ValueError('Unsupported empty tuple of masks')

    shape = self.masks[0].shape
    for mask in self.masks[1:]:
      if shape != mask.shape:
        raise ValueError(
            f'Unexpected mask shape, got: {mask.shape}, expected: {shape}'
        )

    if not all(isinstance(mask, Mask) for mask in self.masks):
      raise ValueError('masks should be of type Mask')

    if any(isinstance(mask, MultiHeadMask) for mask in self.masks):
      raise ValueError('Nesting MultiHeadMasks is not supported')

  @property
  def shape(self) -> tuple[int, ...]:
    return (len(self.masks),) + self.masks[0].shape

  def __getitem__(self, idx) -> np.ndarray:
    if len(idx) != 3:
      raise NotImplementedError(f'Unsupported slice: {idx}')

    head_slice = idx[0]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Build every head mask with the same (q_seq_len, kv_seq_len) and vary only the boolean pattern
2. Precompute shape once and pass it to every per-head mask constructor
3. Check shapes in a loop before constructing MultiHeadMask

Example fix

// before
MultiHeadMask([make_mask(head_shape(h)) for h in heads])  # shapes differ
// after
shape = (q_len, kv_len)
MultiHeadMask([make_mask(shape, h) for h in heads])

shape = masks[0].shape
assert all(m.shape == shape for m in masks), [m.shape for m in masks]
mh = MultiHeadMask(masks)

When it happens

Trigger: MultiHeadMask([m1, m2, ...]) where some mi.shape differs from masks[0].shape, e.g. per-head masks built with head-specific sequence lengths or a bug in a per-head loop.

Common situations: Building per-head sliding-window masks with different radii implemented as different shapes instead of different values; mixing a causal mask of one length with head-specific masks of another; inconsistent kv cache lengths per head.

Related errors

• Invalid shape for other: {other.shape}, expected: {self.shap
• Masks must have the same shape
• Nesting MultiHeadMasks is not supported
• Mask function must return a boolean-valued array, but got: {
• partial_mask_blocks must be of type np.bool_ but got {partia

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