{"record":{"id":"ac948d6810ebc228","repo":"xai-org/x-algorithm","slug":"mask-dimensionality-mask-ndim-must-match-logits","errorCode":null,"errorMessage":"Mask dimensionality {mask.ndim} must match logits dimensionality {attn_logits.ndim} for {mask.shape}/{attn_logits.shape}.","messagePattern":"Mask dimensionality (.+?) must match logits dimensionality (.+?) for (.+?)/(.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/models/attention.py","lineNumber":308,"sourceCode":"            )\n            segment_ids_for_keys = segment_ids_k if segment_ids_k is not None else segment_ids\n            if segment_ids_k is not None:\n                segment_ids_for_keys = with_sharding_constraint(\n                    segment_ids_for_keys,\n                    sharding_rule(\n                        NamedShape(segment_ids_for_keys.shape, (\"batch_attn\", \"replicated\")),\n                    ),\n                )\n            segment_mask = make_attention_mask(\n                segment_ids, segment_ids_for_keys, jnp.equal, dtype=query.dtype\n            )\n            mask *= segment_mask\n\n        mask = mask[:, :, None, :, :]\n\n        if mask is not None:\n            if mask.ndim != attn_logits.ndim:\n                raise ValueError(\n                    f\"Mask dimensionality {mask.ndim} must match logits dimensionality \"\n                    f\"{attn_logits.ndim} for {mask.shape}/{attn_logits.shape}.\"\n                )\n            attn_logits = jnp.where(mask, attn_logits, -1e30)\n\n        attn_weights = jax.nn.softmax(attn_logits).astype(query.dtype)\n\n        attn = jnp.einsum(\"...hHtT,...Thd->...thHd\", attn_weights, value)\n        attn = with_sharding_constraint(\n            attn,\n            sharding_rule(\n                NamedShape(\n                    attn.shape, (\"batch_attn\", \"replicated\", \"head\", \"replicated\", \"hidden\")\n                ),\n            ),\n        )\n\n        return attn","sourceCodeStart":290,"sourceCodeEnd":326,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/models/attention.py#L290-L326","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Reshape the mask to logits rank, typically mask[:, None, None, :] broadcast or a [B, S, S] pattern expanded per call_attn's expectations.","Print mask.shape and attn_logits.shape at the failure point and align dimensions (batch, heads, q-len, k-len).","Let the layer build the mask internally (drop custom masks from extra_attn_kwargs) if you only need padding/segment masking."],"exampleFix":"# before\nextra_attn_kwargs[\"masks\"] = mask  # shape [B, S]\n\n# after\nextra_attn_kwargs[\"masks\"] = mask[:, None, None, :]  # broadcast to [B, H, S, S]","handlingStrategy":"validation","validationCode":"expected_logits_ndim = 4  # [B, H, S, S]\nassert mask.ndim + 1 == expected_logits_ndim or mask.ndim == expected_logits_ndim","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Log mask.shape alongside attn logits shape in debug builds.","Build masks via the layer's own helpers instead of hand-crafted tensors when switching attn_impl."],"tags":["attention","mask-shape","broadcasting","jax"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}