jax-ml/jax · error · ValueError

Unexpected mask shape: {mask.shape}

Error message

Unexpected mask shape: {mask.shape}

What it means

_make_splash_attention expects a rank-3 mask array shaped (num_heads, q_seq_len, kv_seq_len) (or an equivalent Mask object). A mask of any other rank is rejected before kernel construction.

Source

Thrown at jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py:2560

    )


def _make_splash_attention(
    mask: np.ndarray | jax.Array | mask_lib.MultiHeadMask,
    *,
    block_sizes: BlockSizes | None = None,
    is_mqa: bool,
    save_residuals: bool = False,
    mask_value: float = DEFAULT_MASK_VALUE,
    attn_logits_soft_cap: float | None = None,
    downcast_smem_data: bool = True,
    head_shards: int,
    q_seq_shards: int,
    residual_checkpoint_name: str | None = None,
    interpret: bool = False,
):
  if len(mask.shape) != 3:
    raise ValueError(f'Unexpected mask shape: {mask.shape}')

  if isinstance(mask, np.ndarray):
    mask = mask_lib.MultiHeadMask(
        [mask_lib.NumpyMask(head_mask) for head_mask in mask]
    )

  if block_sizes is None:
    block_sizes = BlockSizes.get_default()

  process_mask_fn = (
      mask_info_lib.process_dynamic_mask
      if isinstance(mask, jax.Array)
      else mask_info_lib.process_mask
  )

  process_mask_dvk_fn = (
      mask_info_lib.process_dynamic_mask_dkv
      if isinstance(mask, jax.Array)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshape/broadcast the mask to (num_heads, q_seq_len, kv_seq_len)
2. Wrap a per-head 2D mask: mask_lib.MultiHeadMask([NumpyMask(m) for m in masks])
3. For a shared causal mask, stack it num_heads times along axis 0

Example fix

// before
mask = make_causal_mask((q_len, kv_len))  # rank 2
// after
mask = np.tile(make_causal_mask((q_len, kv_len)), (num_heads, 1, 1))

assert isinstance(mask, Mask) or (isinstance(mask, np.ndarray) and mask.ndim == 3), mask.shape

def is_splash_mask(m):
    return hasattr(m, 'shape') or (isinstance(m, np.ndarray) and m.ndim == 3)

When it happens

Trigger: Passing a 2D (single-head) numpy mask, a 4D batch mask, or a scalar mask to make_splash_attention.

Common situations: Migrating code from single-head attention that built (q, kv) masks; passing batched masks shaped (batch, heads, q, kv); forgetting to broadcast a causal mask across heads.

Related errors

• Mask function must return a boolean-valued array, but got: {
• partial_mask_blocks must be of type np.bool_ but got {partia
• Invalid shape for other: {other.shape}, expected: {self.shap
• chunk_size must be positive
• Masks must have the same shape

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