{"record":{"id":"bf36d5bbff69a230","repo":"jax-ml/jax","slug":"masks-must-have-the-same-shape","errorCode":null,"errorMessage":"Masks must have the same shape","messagePattern":"Masks must have the same shape","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask.py","lineNumber":138,"sourceCode":"  return mask\n\n\ndef make_random_mask(\n    shape: tuple[int, int], sparsity: float, seed: int\n) -> np.ndarray:\n  \"\"\"Makes a random attention mask.\"\"\"\n  np.random.seed(seed)\n  return np.random.binomial(n=1, p=1.0 - sparsity, size=shape).astype(np.bool_)\n\n\n@dataclasses.dataclass\nclass LogicalOr(Mask):\n  left: Mask\n  right: Mask\n\n  def __init__(self, left: Mask, right: Mask):\n    if left.shape != right.shape:\n      raise ValueError('Masks must have the same shape')\n    self.left = left\n    self.right = right\n\n  @property\n  def shape(self) -> tuple[int, ...]:\n    return self.left.shape\n\n  def __getitem__(self, idx) -> np.ndarray:\n    return self.left[idx] | self.right[idx]\n\n  def __hash__(self):\n    return hash((type(self),) + (self.left, self.right))\n\n\n@dataclasses.dataclass\nclass LogicalAnd(Mask):\n  left: Mask\n  right: Mask","sourceCodeStart":120,"sourceCodeEnd":156,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask.py#L120-L156","documentation":"LogicalOr (built via Mask.__or__) re-validates at construction that its two operand masks have equal shapes. Direct instantiation LogicalOr(left, right) with mismatched shapes raises this, mirroring the __or__ check.","triggerScenarios":"Constructing LogicalOr(left, right) directly, or via '|', where left.shape != right.shape; also hit when one operand's shape changes after being wrapped lazily.","commonSituations":"Building composite mask expression trees programmatically and mixing masks of different sequence lengths; lazy masks whose shape depends on mutable state.","solutions":["Use '|' on masks already verified to share (q, kv) shapes","Validate shapes at mask-construction time in your factory function","Avoid constructing LogicalOr/LogicalAnd directly; rely on operators"],"exampleFix":"// before\ncombined = LogicalOr(causal_1024, local_512)\n// after\nlocal = make_local_mask((1024, 1024), ...)\ncombined = causal_1024 | local","handlingStrategy":"validation","validationCode":"assert left.shape == right.shape\ncombined = left | right  # let __or__ build LogicalOr","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Prefer operators over direct LogicalOr construction","Validate shapes in one place before composing"],"tags":["jax","splash-attention","mask","shape-mismatch"],"backgroundTag":"mask-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}