jax-ml/jax · error · ValueError

Nesting MultiHeadMasks is not supported

Error message

Nesting MultiHeadMasks is not supported

What it means

A MultiHeadMask whose elements include another MultiHeadMask is rejected: nesting would create ragged per-head structures that the kernel's per-head lowering cannot handle. Each element must be a flat (q, kv) Mask.

Source

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

  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]
    if isinstance(head_slice, int):
      assert head_slice >= 0 and head_slice <= len(self.masks)
      return self.masks[head_slice][idx[1:]]
    else:
      slice_masks = [mask[idx[1:]] for mask in self.masks[head_slice]]
      return np.stack(slice_masks)

  def __eq__(self, other: object):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Flatten nested MultiHeadMasks into a single list of per-head masks before wrapping
2. Combine same-shape alternatives with '|' or '&' rather than nesting wrappers
3. If all heads share one mask, use a single NumpyMask or rank-3 array instead

Example fix

// before
left = MultiHeadMask([m]*4); right = MultiHeadMask([n]*4)
combined = MultiHeadMask([left, right])
// after
combined = MultiHeadMask([mi | ni for mi, ni in zip([m]*4, [n]*4)])

assert not any(isinstance(m, MultiHeadMask) for m in masks)

def is_flat_mask_list(ms):
    return all(isinstance(m, Mask) and not isinstance(m, MultiHeadMask) for m in ms)

When it happens

Trigger: MultiHeadMask([...]) where any element was produced by combining multiple masks via MultiHeadMask instead of '|'; passing a MultiHeadMask inside the masks list during programmatic construction.

Common situations: Trying to express grouped-query or grouped-head masks by nesting; recursive mask-building utilities that wrap their result in MultiHeadMask at every level; merging per-group masks into a multi-head wrapper.

Related errors

• Unexpected mask shape, got: {mask.shape}, expected: {shape}
• Mask function must return a boolean-valued array, but got: {
• partial_mask_blocks must be of type np.bool_ but got {partia
• Unexpected mask shape: {mask.shape}
• Invalid shape for other: {other.shape}, expected: {self.shap

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