jax-ml/jax · error · ValueError

masks should be of type Mask

Error message

masks should be of type Mask

What it means

MultiHeadMask requires every element of its masks sequence to be an instance of the Mask base class. Raw numpy arrays, callables, or other objects are rejected because the kernel lowering relies on the lazy Mask protocol.

Source

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

@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]
    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]]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap arrays: MultiHeadMask([mask_lib.NumpyMask(a) for a in arrays])
2. Or pass a single rank-3 numpy mask directly to make_splash_attention, which auto-wraps it
3. Ensure all mask builders return Mask instances, not bare arrays

Example fix

// before
MultiHeadMask([make_causal_mask(shape) for _ in range(h)])  # ndarrays
// after
MultiHeadMask([mask_lib.NumpyMask(make_causal_mask(shape)) for _ in range(h)])

assert all(isinstance(m, Mask) for m in masks)

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

When it happens

Trigger: MultiHeadMask([arr1, arr2]) where elements are np.ndarray instead of Mask objects; passing a generator or list of arrays from make_causal_mask (which returns ndarrays).

Common situations: Wrapping raw causal numpy masks per head without converting to NumpyMask; refactoring mask code from arrays to lazy masks halfway; dataclass field typed loosely accepting any sequence.

Related errors

• 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
• chunk_size must be positive

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