{"record":{"id":"72e26930751c71a7","repo":"jax-ml/jax","slug":"masks-should-be-of-type-mask","errorCode":null,"errorMessage":"masks should be of type Mask","messagePattern":"masks should be of type Mask","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask.py","lineNumber":193,"sourceCode":"@dataclasses.dataclass\nclass MultiHeadMask(Mask):\n  \"\"\"Lazy multihead mask, combines multiple lazy masks one per head.\"\"\"\n\n  masks: Sequence[Mask]\n\n  def __post_init__(self):\n    if not self.masks:\n      raise ValueError('Unsupported empty tuple of masks')\n\n    shape = self.masks[0].shape\n    for mask in self.masks[1:]:\n      if shape != mask.shape:\n        raise ValueError(\n            f'Unexpected mask shape, got: {mask.shape}, expected: {shape}'\n        )\n\n    if not all(isinstance(mask, Mask) for mask in self.masks):\n      raise ValueError('masks should be of type Mask')\n\n    if any(isinstance(mask, MultiHeadMask) for mask in self.masks):\n      raise ValueError('Nesting MultiHeadMasks is not supported')\n\n  @property\n  def shape(self) -> tuple[int, ...]:\n    return (len(self.masks),) + self.masks[0].shape\n\n  def __getitem__(self, idx) -> np.ndarray:\n    if len(idx) != 3:\n      raise NotImplementedError(f'Unsupported slice: {idx}')\n\n    head_slice = idx[0]\n    if isinstance(head_slice, int):\n      assert head_slice >= 0 and head_slice <= len(self.masks)\n      return self.masks[head_slice][idx[1:]]\n    else:\n      slice_masks = [mask[idx[1:]] for mask in self.masks[head_slice]]","sourceCodeStart":175,"sourceCodeEnd":211,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask.py#L175-L211","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Wrap arrays: MultiHeadMask([mask_lib.NumpyMask(a) for a in arrays])","Or pass a single rank-3 numpy mask directly to make_splash_attention, which auto-wraps it","Ensure all mask builders return Mask instances, not bare arrays"],"exampleFix":"// before\nMultiHeadMask([make_causal_mask(shape) for _ in range(h)])  # ndarrays\n// after\nMultiHeadMask([mask_lib.NumpyMask(make_causal_mask(shape)) for _ in range(h)])","handlingStrategy":"type-guard","validationCode":"assert all(isinstance(m, Mask) for m in masks)","typeGuard":"def all_are_masks(ms):\n    return all(isinstance(m, Mask) and not isinstance(m, MultiHeadMask) for m in ms)","tryCatchPattern":null,"preventionTips":["Wrap arrays with NumpyMask before MultiHeadMask","Make mask builders return Mask, not ndarray"],"tags":["jax","splash-attention","mask","type-validation"],"backgroundTag":"type-validation-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}