xai-org/grok-1 · 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

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.

Source

Thrown at model.py:871

        query_heads = jnp.reshape(query_heads, (b, t, kv_h, h // kv_h, d))
        query_heads = with_sharding_constraint(
            query_heads, P(self.data_axis, None, "model", None, None)
        )

        # Compute attention weights.
        # Attention softmax is always carried out in fp32.
        attn_logits = jnp.einsum("...thHd,...Thd->...hHtT", query_heads, key_heads).astype(
            jnp.float32
        )
        attn_logits *= self.attn_output_multiplier
        max_attn_val = jnp.array(30.0, dtype=attn_logits.dtype)
        attn_logits = max_attn_val * jnp.tanh(attn_logits / max_attn_val)

        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)  # [H, T', T]

        # Weight the values by the attention and flatten the head vectors.
        attn = jnp.einsum("...hHtT,...Thd->...thHd", attn_weights, value_heads)
        attn = with_sharding_constraint(attn, P(self.data_axis, None, "model", None, None))
        leading_dims = attn.shape[:2]
        attn = jnp.reshape(attn, (*leading_dims, -1))  # [T', H*V]
        attn = with_sharding_constraint(attn, P(self.data_axis, None, "model"))
        # Apply another projection to get the final embeddings.
        final_projection = Linear(
            self.model_size,
            with_bias=False,
            sharding=P("model", "data"),
            mesh=mesh,

View on GitHub (pinned to 7050ed204b)

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.

Example fix

# before: [B, T, T] causal mask (3-d) -> after expand 4-d != 5-d logits
mask = jnp.tril(jnp.ones((B, T, T), dtype=bool))

# after: add the group axis so the expand makes it 5-d
mask = jnp.tril(jnp.ones((B, T, T), dtype=bool))[:, :, None, :]   # [B, T, 1, T, 1]
# code's mask[:, :, None, :, :] then gives 5-d, matching [..., h, H, t, T] logits

Defensive patterns

Strategy: validation

Validate before calling

import jax.numpy as jnp

LOGITS_NDIM = 5  # [..., t(h-group), H, t', T] from '...thHd,...Thd->...hHtT'

def check_attention_mask(mask: jax.Array) -> jax.Array:
    if mask is None:
        raise TypeError('mask=None is unsupported: indexing runs before the None check')
    # model.py inserts one axis via mask[:, :, None, :, :], so pre-expand rank must be 4
    if mask.ndim + 1 != LOGITS_NDIM:
        mask = mask.reshape(mask.shape[0], mask.shape[-2], 1, mask.shape[-1], 1)  # adapt as needed
    assert mask.ndim + 1 == LOGITS_NDIM, f'mask ndim {mask.ndim} vs logits ndim {LOGITS_NDIM}'
    return mask

Type guard

def mask_rank_ok(mask: jax.Array) -> bool:
    """Mask is 4-d so the built-in axis expansion yields 5-d, matching logits."""
    return mask is not None and mask.ndim == 4

Try / catch

try:
    out = model.apply(params, tokens, mask=mask)
except ValueError as e:
    if 'Mask dimensionality' in str(e):
        mask = mask[:, None, :, :] if mask.ndim == 4 else mask[:, :, None, :, :]
        # only retry once with corrected rank; re-raise if still failing
        out = model.apply(params, tokens, mask=mask)
    else:
        raise

Prevention

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.

When it happens

Trigger: 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.

Common situations: 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.

Related errors

AI-assisted analysis of xai-org/grok-1@7050ed204b (2026-08-15). Data as JSON: /api/errors/4a51bc3f3731815c. Report an issue: GitHub.