jax-ml/jax · error · ValueError

Expected a 3-dim mask, instead got: {mask.shape=}

Error message

Expected a 3-dim mask, instead got: {mask.shape=}

What it means

_process_mask expects the attention mask to be a 3D array shaped (num_heads, q_seq_len, kv_seq_len). A 2D mask (batch, seq) or 4D mask (batch, heads, q, kv) fails this check immediately.

Source

Thrown at jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask_info.py:561

      a data type smaller than np.int32 (if possible).
    head_shards: Number of head shards of the mesh in which the kernel is
      launched.
    q_seq_shards: Number of Q sequence shards of the mesh in which the kernel is
      launched.
    shrink_grid: Whether or not we should apply the grid shrinking optimization.

  Returns:
    `MaskInfo`, a sparse representation of the dense mask.
    `MaskCallable`: a callable that, given in input Q and KV indices, returns
      the value of the mask at those coordinates.

  Raises:
    ValueError: if the input mask is invalid or the block sizes are not
    compatible with the mask sizes.
  """

  if len(mask.shape) != 3:
    raise ValueError(f'Expected a 3-dim mask, instead got: {mask.shape=}')

  head_count, q_seq_len, kv_seq_len = mask.shape
  q_block_size, kv_block_size = block_shape
  q_blocks_count, q_mod = divmod(q_seq_len, q_block_size)
  kv_blocks_count, kv_mod = divmod(kv_seq_len, kv_block_size)

  if q_mod != 0:
    raise ValueError(f'{q_block_size=} should divide {q_seq_len=}.')
  if kv_mod != 0:
    raise ValueError(f'{kv_block_size=} should divide {kv_seq_len=}.')

  q_seq_len_per_shard, mod = divmod(q_seq_len, q_seq_shards)
  if mod != 0:
    raise ValueError(f'{q_seq_shards=} should divide {q_seq_len=}.')

  q_blocks_per_shard, mod = divmod(q_seq_len_per_shard, q_block_size)
  if mod != 0:
    raise ValueError(f'{q_block_size=} should divide {q_seq_len_per_shard=}.')

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Expand 2D masks to 3D: mask[None, :, :] repeated/tiled across heads
2. Squeeze the batch dimension from 4D masks (splash attention handles batching via scan/sharding, not a batch dim)
3. Tile per-head masks with jnp.tile(mask[None], (num_heads, 1, 1))

Example fix

# before
mask_2d = tokens != 0  # (q, kv)
# after
mask_3d = jnp.tile(mask_2d[None], (num_heads, 1, 1))  # (heads, q, kv)

assert mask.ndim == 3, mask.shape

def is_3d_mask(m) -> bool:
    return getattr(m, 'ndim', -1) == 3

When it happens

Trigger: Passing a standard Bert-style 2D padding mask or a 4D transformer mask directly to make_splash_attention_mask instead of the expected 3D (head, q, kv) layout.

Common situations: Porting a HuggingFace model whose attention_mask is (batch, seq); broadcasting a per-token mask without adding the head dimension first.

Related errors

• Mask function must return a boolean-valued array, but got: {
• partial_mask_blocks must be of type np.bool_ but got {partia
• Out-of-bounds masked swap of ({device_id} {local_core_id} {m
• Higher-order AD not supported.
• Block sizes for dq kernel are not needed with a fused kernel

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