{"record":{"id":"c36212abb1c58370","repo":"jax-ml/jax","slug":"invalid-shape-for-other-other-shape-expected","errorCode":null,"errorMessage":"Invalid shape for other: {other.shape}, expected: {self.shape}","messagePattern":"Invalid shape for other: (.+?), expected: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask.py","lineNumber":43,"sourceCode":"class Mask:\n  \"\"\"A base class for splash attention masks.\"\"\"\n\n  @property\n  def shape(self) -> tuple[int, ...]:\n    raise NotImplementedError\n\n  def __getitem__(self, idx) -> np.ndarray:\n    raise NotImplementedError\n\n  def __bool__(self) -> bool:\n    raise NotImplementedError(\n        'Conversion to bool is unsupported. Could be caused by using logical'\n        ' instead of bitwise operations on masks.'\n    )\n\n  def __or__(self, other: Mask) -> Mask:\n    if self.shape != other.shape:\n      raise ValueError(\n          f'Invalid shape for other: {other.shape}, expected: {self.shape}'\n      )\n    return LogicalOr(self, other)\n\n  def __and__(self, other: Mask) -> Mask:\n    if self.shape != other.shape:\n      raise ValueError(\n          f'Invalid shape for other: {other.shape}, expected: {self.shape}'\n      )\n    return LogicalAnd(self, other)\n\n\ndef make_causal_mask(shape: tuple[int, int], offset: int = 0) -> np.ndarray:\n  \"\"\"Makes a causal attention mask.\n\n  Args:\n    shape: Shape of the 2-dim mask: (q_seq_len, kv_seq_len).\n    offset: Offset of q start wrt kv. A positive offset shifts the bottom","sourceCodeStart":25,"sourceCodeEnd":61,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask.py#L25-L61","documentation":"Mask.__or__ builds a lazy LogicalOr from two masks that must have identical shapes; combining masks of different (q, kv) shapes is rejected because there is no defined union of differently-shaped regions.","triggerScenarios":"mask_a | mask_b where mask_a.shape != mask_b.shape, e.g. a causal mask for seq_len 1024 unioned with a padding mask for seq_len 512.","commonSituations":"Sequence lengths changing between mask construction and composition (e.g. after padding/truncation); masks built from different shape constants; off-by-one kv lengths for cache prefix.","solutions":["Rebuild both masks with the same (q_seq_len, kv_seq_len) shape before combining","Pad or crop one mask to match the other's shape","Centralize shape constants so all masks derive from one (q_len, kv_len) source"],"exampleFix":"// before\ncombined = make_causal_mask((1024,1024)) | padding_mask_512\n// after\ncombined = make_causal_mask((512,512)) | padding_mask_512","handlingStrategy":"validation","validationCode":"assert a.shape == b.shape, (a.shape, b.shape)\ncombined = a | b","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Build all masks from one (q_len, kv_len) tuple","Rebuild masks whenever sequence length changes"],"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"}