jax-ml/jax · error · ValueError

Masks must have the same shape

Error message

Masks must have the same shape

What it means

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.

Source

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

  return mask


def make_random_mask(
    shape: tuple[int, int], sparsity: float, seed: int
) -> np.ndarray:
  """Makes a random attention mask."""
  np.random.seed(seed)
  return np.random.binomial(n=1, p=1.0 - sparsity, size=shape).astype(np.bool_)


@dataclasses.dataclass
class LogicalOr(Mask):
  left: Mask
  right: Mask

  def __init__(self, left: Mask, right: Mask):
    if left.shape != right.shape:
      raise ValueError('Masks must have the same shape')
    self.left = left
    self.right = right

  @property
  def shape(self) -> tuple[int, ...]:
    return self.left.shape

  def __getitem__(self, idx) -> np.ndarray:
    return self.left[idx] | self.right[idx]

  def __hash__(self):
    return hash((type(self),) + (self.left, self.right))


@dataclasses.dataclass
class LogicalAnd(Mask):
  left: Mask
  right: Mask

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use '|' on masks already verified to share (q, kv) shapes
2. Validate shapes at mask-construction time in your factory function
3. Avoid constructing LogicalOr/LogicalAnd directly; rely on operators

Example fix

// before
combined = LogicalOr(causal_1024, local_512)
// after
local = make_local_mask((1024, 1024), ...)
combined = causal_1024 | local

assert left.shape == right.shape
combined = left | right  # let __or__ build LogicalOr

When it happens

Trigger: Constructing LogicalOr(left, right) directly, or via '|', where left.shape != right.shape; also hit when one operand's shape changes after being wrapped lazily.

Common situations: Building composite mask expression trees programmatically and mixing masks of different sequence lengths; lazy masks whose shape depends on mutable state.

Related errors

• Invalid shape for other: {other.shape}, expected: {self.shap
• Unexpected mask shape, got: {mask.shape}, expected: {shape}
• 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}

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