jax-ml/jax · error · ValueError

Invalid shape for other: {other.shape}, expected: {self.shap

Error message

Invalid shape for other: {other.shape}, expected: {self.shape}

What it means

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.

Source

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

class Mask:
  """A base class for splash attention masks."""

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

  def __getitem__(self, idx) -> np.ndarray:
    raise NotImplementedError

  def __bool__(self) -> bool:
    raise NotImplementedError(
        'Conversion to bool is unsupported. Could be caused by using logical'
        ' instead of bitwise operations on masks.'
    )

  def __or__(self, other: Mask) -> Mask:
    if self.shape != other.shape:
      raise ValueError(
          f'Invalid shape for other: {other.shape}, expected: {self.shape}'
      )
    return LogicalOr(self, other)

  def __and__(self, other: Mask) -> Mask:
    if self.shape != other.shape:
      raise ValueError(
          f'Invalid shape for other: {other.shape}, expected: {self.shape}'
      )
    return LogicalAnd(self, other)


def make_causal_mask(shape: tuple[int, int], offset: int = 0) -> np.ndarray:
  """Makes a causal attention mask.

  Args:
    shape: Shape of the 2-dim mask: (q_seq_len, kv_seq_len).
    offset: Offset of q start wrt kv. A positive offset shifts the bottom

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Rebuild both masks with the same (q_seq_len, kv_seq_len) shape before combining
2. Pad or crop one mask to match the other's shape
3. Centralize shape constants so all masks derive from one (q_len, kv_len) source

Example fix

// before
combined = make_causal_mask((1024,1024)) | padding_mask_512
// after
combined = make_causal_mask((512,512)) | padding_mask_512

assert a.shape == b.shape, (a.shape, b.shape)
combined = a | b

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Masks must have the same shape
• 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/c36212abb1c58370. Report an issue: GitHub.