jax-ml/jax · error · ValueError

Unsupported empty tuple of masks

Error message

Unsupported empty tuple of masks

What it means

MultiHeadMask wraps one Mask per attention head; an empty sequence has no head and no shape, so construction with masks=[] is rejected in __post_init__.

Source

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

  def shape(self) -> tuple[int, ...]:
    return self.left.shape

  def __getitem__(self, idx) -> np.ndarray:
    return self.left[idx] & self.right[idx]

  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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Guard against empty lists: MultiHeadMask(masks) only if masks, else use a shared single mask
2. Validate num_heads > 0 before building the mask
3. Default to a shared causal mask (rank-3 numpy array) when there are no per-head masks

Example fix

// before
MultiHeadMask([make_head_mask(h, shape) for h in heads if keep(h)])  # may be []
// after
masks = [make_head_mask(h, shape) for h in heads if keep(h)]
assert masks, 'no head masks'
MultiHeadMask(masks)

assert masks, 'cannot build MultiHeadMask with zero masks'
mh = MultiHeadMask(masks)

When it happens

Trigger: Creating MultiHeadMask([]) or MultiHeadMask(masks) where masks is an empty list, e.g. from a list comprehension over zero heads or a config with num_heads=0.

Common situations: num_heads read as 0 from config; iterating head configs that come back empty after filtering; dynamic construction where the head list is populated later.

Related errors

• other requires mask to be provided
• mask cannot be a block if pointer is not a block
• Mask function must return a boolean-valued array, but got: {
• partial_mask_blocks must be of type np.bool_ but got {partia
• {bkv_compute=} should not be greater than {bkv=}

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