{"record":{"id":"764f85da789c103e","repo":"jax-ml/jax","slug":"unexpected-mask-shape-got-mask-shape-expected","errorCode":null,"errorMessage":"Unexpected mask shape, got: {mask.shape}, expected: {shape}","messagePattern":"Unexpected mask shape, got: (.+?), expected: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask.py","lineNumber":188,"sourceCode":"\n  def __hash__(self):\n    return hash((type(self),) + (self.left, self.right))\n\n\n@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]","sourceCodeStart":170,"sourceCodeEnd":206,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask.py#L170-L206","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Build every head mask with the same (q_seq_len, kv_seq_len) and vary only the boolean pattern","Precompute shape once and pass it to every per-head mask constructor","Check shapes in a loop before constructing MultiHeadMask"],"exampleFix":"// before\nMultiHeadMask([make_mask(head_shape(h)) for h in heads])  # shapes differ\n// after\nshape = (q_len, kv_len)\nMultiHeadMask([make_mask(shape, h) for h in heads])","handlingStrategy":"validation","validationCode":"shape = masks[0].shape\nassert all(m.shape == shape for m in masks), [m.shape for m in masks]\nmh = MultiHeadMask(masks)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Pass one shape to all per-head mask builders","Vary mask values, not shapes, across heads"],"tags":["jax","splash-attention","mask","shape-mismatch","multi-head"],"backgroundTag":"mask-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}