xai-org/x-algorithm · error · ValueError

Mask dimensionality {mask.ndim} must match logits dimensiona

Error message

Mask dimensionality {mask.ndim} must match logits dimensionality {attn_logits.ndim} for {mask.shape}/{attn_logits.shape}.

What it means

In JaxAttention.call_attn the (already expanded with mask[:, :, None, :, :]) mask must have the same ndim as the attention logits so jnp.where(mask, logits, -1e30) broadcasts cleanly. A mismatch (e.g. a 3D or 4D mask against 4D logits) means the mask was shaped for a different attention layout and the error reports both ndims and shapes.

Source

Thrown at phoenix/xrex/models/attention.py:308

            )
            segment_ids_for_keys = segment_ids_k if segment_ids_k is not None else segment_ids
            if segment_ids_k is not None:
                segment_ids_for_keys = with_sharding_constraint(
                    segment_ids_for_keys,
                    sharding_rule(
                        NamedShape(segment_ids_for_keys.shape, ("batch_attn", "replicated")),
                    ),
                )
            segment_mask = make_attention_mask(
                segment_ids, segment_ids_for_keys, jnp.equal, dtype=query.dtype
            )
            mask *= segment_mask

        mask = mask[:, :, None, :, :]

        if mask is not None:
            if mask.ndim != attn_logits.ndim:
                raise ValueError(
                    f"Mask dimensionality {mask.ndim} must match logits dimensionality "
                    f"{attn_logits.ndim} for {mask.shape}/{attn_logits.shape}."
                )
            attn_logits = jnp.where(mask, attn_logits, -1e30)

        attn_weights = jax.nn.softmax(attn_logits).astype(query.dtype)

        attn = jnp.einsum("...hHtT,...Thd->...thHd", attn_weights, value)
        attn = with_sharding_constraint(
            attn,
            sharding_rule(
                NamedShape(
                    attn.shape, ("batch_attn", "replicated", "head", "replicated", "hidden")
                ),
            ),
        )

        return attn

View on GitHub (pinned to 24c60942c5)

Solutions

1. Reshape the mask to logits rank, typically mask[:, None, None, :] broadcast or a [B, S, S] pattern expanded per call_attn's expectations.
2. Print mask.shape and attn_logits.shape at the failure point and align dimensions (batch, heads, q-len, k-len).
3. Let the layer build the mask internally (drop custom masks from extra_attn_kwargs) if you only need padding/segment masking.

Example fix

# before
extra_attn_kwargs["masks"] = mask  # shape [B, S]

# after
extra_attn_kwargs["masks"] = mask[:, None, None, :]  # broadcast to [B, H, S, S]

expected_logits_ndim = 4  # [B, H, S, S]
assert mask.ndim + 1 == expected_logits_ndim or mask.ndim == expected_logits_ndim

When it happens

Trigger: Passing a padding/segment mask of the wrong rank (e.g. [B, S] instead of [B, 1, S, S] or [B, H, S, S]) to jax_attn with masks in extra_attn_kwargs; using a mask batched per-head when the impl expects a shared-head mask or vice versa.

Common situations: Switching attn_impl from pallas/flash (which accept 2D masks) to jax_attn without reshaping; changing num_heads or reshape_layers so the logits gain an axis; masks built for a different sequence layout.

Understand the failure class

Background: Tensor shape mismatch errors ("must have shape", "expected shape ... got ..."): when tensor dimensions disagree with what an op or layer was told to expect — this error's family across 6 libraries.

Related errors

• Please override this method for specific attention impl.
• Invalid backward pass implementation: {backward_pass_impl}
• cap_method must be in [tanh, soft_sign, none], got {cap_meth
• q, k, and v should all be 4D, got: {q.ndim=}, {k.ndim=}, {v.
• Expected {k.shape=} to be {kv_shape} (inferred from q)

AI-assisted analysis of xai-org/x-algorithm@24c60942c5 (2026-08-28). Data as JSON: /api/errors/ac948d6810ebc228. Report an issue: GitHub.