{"record":{"id":"4a51bc3f3731815c","repo":"xai-org/grok-1","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":"model.py","lineNumber":871,"sourceCode":"        query_heads = jnp.reshape(query_heads, (b, t, kv_h, h // kv_h, d))\n        query_heads = with_sharding_constraint(\n            query_heads, P(self.data_axis, None, \"model\", None, None)\n        )\n\n        # Compute attention weights.\n        # Attention softmax is always carried out in fp32.\n        attn_logits = jnp.einsum(\"...thHd,...Thd->...hHtT\", query_heads, key_heads).astype(\n            jnp.float32\n        )\n        attn_logits *= self.attn_output_multiplier\n        max_attn_val = jnp.array(30.0, dtype=attn_logits.dtype)\n        attn_logits = max_attn_val * jnp.tanh(attn_logits / max_attn_val)\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        attn_weights = jax.nn.softmax(attn_logits).astype(query.dtype)  # [H, T', T]\n\n        # Weight the values by the attention and flatten the head vectors.\n        attn = jnp.einsum(\"...hHtT,...Thd->...thHd\", attn_weights, value_heads)\n        attn = with_sharding_constraint(attn, P(self.data_axis, None, \"model\", None, None))\n        leading_dims = attn.shape[:2]\n        attn = jnp.reshape(attn, (*leading_dims, -1))  # [T', H*V]\n        attn = with_sharding_constraint(attn, P(self.data_axis, None, \"model\"))\n        # Apply another projection to get the final embeddings.\n        final_projection = Linear(\n            self.model_size,\n            with_bias=False,\n            sharding=P(\"model\", \"data\"),\n            mesh=mesh,","sourceCodeStart":853,"sourceCodeEnd":889,"githubUrl":"https://github.com/xai-org/grok-1/blob/7050ed204b8206bb8645c7b7bbef7252f79561b0/model.py#L853-L889","documentation":"Raised in Grok-1's attention (model.py:871) when the boolean attention mask's rank does not equal the rank of attn_logits after the mask has been expanded with mask[:, :, None, :, :]. The einsum '...thHd,...Thd->...hHtT' produces 5-d logits [batch, seq', heads_per_group(window/shard), group, seq], so after the None insertion the mask must also be 5-d ([batch, seq', seq, ...]-compatible). This check is the library telling you the mask you injected via the hk.multi_transform or by patching _attention does not have the layout the kernel expects.","triggerScenarios":"Building your own attention mask and passing it into the model path that reaches this block: a plain [batch, seq, seq] (3-d) causal mask — after [:, :, None, :, :] it becomes 4-d while logits are 5-d; a 4-d mask intended 'as-is' (the code inserts the axis for you); or code from an older Grok-1 revision where make_attention_mask returned a different rank being reused with this revision. Note also the latent bug in this region: mask[:, :, None, :, :] runs BEFORE the `if mask is not None` check, so passing None crashes earlier with a TypeError.","commonSituations":"Adding bidirectional/padding masks for fine-tuning and assuming a [B, T, T] or [B, H, T, T] mask like BERT/PyTorch; porting masks between the two attention implementations in this repo (full vs sliding-window shard_map path have different expected ranks); upgrading the xai-org/grok-1 checkout and carrying a local _attention patch forward without re-checking rank.","solutions":["Construct the mask with the repo's own helper: causal_mask = make_attention_mask(logits=..., scores=jnp.zeros([1, 1, 1, 1, 1])) or reuse exactly the mask shape built inside model.py's attention, so ranks stay in sync.","If building manually, produce a 4-d mask of shape [batch, q_len, kv_len, 1] (or matching [B, T', T, G] pattern); the code's [:, :, None, :, :] then yields the required 5-d [B, T', 1, G, T]. Concretely: mask = mask[:, None, None, :, :] style bookkeeping — print mask.ndim and attn_logits.ndim side by side until equal.","Diff your local model.py against the upstream xai-org/grok-1 model.py (git diff model.py) to catch stale mask plumbing from an older revision.","Never pass None here despite the `if mask is not None` guard — the indexing above it raises first; use an all-True mask of the correct shape instead."],"exampleFix":"# before: [B, T, T] causal mask (3-d) -> after expand 4-d != 5-d logits\nmask = jnp.tril(jnp.ones((B, T, T), dtype=bool))\n\n# after: add the group axis so the expand makes it 5-d\nmask = jnp.tril(jnp.ones((B, T, T), dtype=bool))[:, :, None, :]   # [B, T, 1, T, 1]\n# code's mask[:, :, None, :, :] then gives 5-d, matching [..., h, H, t, T] logits","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\n\nLOGITS_NDIM = 5  # [..., t(h-group), H, t', T] from '...thHd,...Thd->...hHtT'\n\ndef check_attention_mask(mask: jax.Array) -> jax.Array:\n    if mask is None:\n        raise TypeError('mask=None is unsupported: indexing runs before the None check')\n    # model.py inserts one axis via mask[:, :, None, :, :], so pre-expand rank must be 4\n    if mask.ndim + 1 != LOGITS_NDIM:\n        mask = mask.reshape(mask.shape[0], mask.shape[-2], 1, mask.shape[-1], 1)  # adapt as needed\n    assert mask.ndim + 1 == LOGITS_NDIM, f'mask ndim {mask.ndim} vs logits ndim {LOGITS_NDIM}'\n    return mask","typeGuard":"def mask_rank_ok(mask: jax.Array) -> bool:\n    \"\"\"Mask is 4-d so the built-in axis expansion yields 5-d, matching logits.\"\"\"\n    return mask is not None and mask.ndim == 4","tryCatchPattern":"try:\n    out = model.apply(params, tokens, mask=mask)\nexcept ValueError as e:\n    if 'Mask dimensionality' in str(e):\n        mask = mask[:, None, :, :] if mask.ndim == 4 else mask[:, :, None, :, :]\n        # only retry once with corrected rank; re-raise if still failing\n        out = model.apply(params, tokens, mask=mask)\n    else:\n        raise","preventionTips":["Reuse model.py's own mask builders (make_attention_mask / the causal mask construction in the forward fn) instead of hand-rolling ranks.","Add an assert mask.ndim == 4 immediately after building any custom mask.","Remember this code inserts one axis itself; do not pre-expand to 5-d 'to be safe'.","Never pass None as the mask here; the indexing before the None check raises a TypeError first."],"tags":["grok-1","jax","attention","mask","shape-error","einsum"],"backgroundTag":null,"analyzedSha":"7050ed204b8206bb8645c7b7bbef7252f79561b0","analyzedAt":"2026-08-15T04:18:25.087Z","schemaVersion":2},"datasetVersion":"2026-08-15T22:17:37.221Z"}