xai-org/grok-1 · error · ValueError

Input must not be scalar.

Error message

Input must not be scalar.

What it means

Raised inside the MoeBlock's hk.transparent _router_weights helper in model.py:259 when the tensor x passed to the MoE router has an empty shape tuple (a JAX scalar, x.shape == ()). The router needs x.shape[-1] to define the input dimension of its 'w' parameter [input_size, num_experts], and a scalar has no trailing axis, so the library bails out with an explicit ValueError instead of producing a confusing IndexError downstream. Any path that funnels a 0-d array into the expert router (e.g. per-token inputs that were over-squeezed) hits this.

Source

Thrown at model.py:259

        routing_logits = self._router_weights(inputs, num_experts, sharding=P("data"))
        assert routing_logits.dtype == jnp.float32
        routing_probs = jax.nn.softmax(routing_logits)

        if padding_mask is not None:
            routing_probs *= padding_mask

        return routing_probs, routing_logits, 0

    @hk.transparent
    def _router_weights(
        self,
        x: jax.Array,
        num_experts: int,
        sharding: Optional[P] = None,
    ):
        fprop_dtype = x.dtype
        if not x.shape:
            raise ValueError("Input must not be scalar.")

        input_size = self.input_size = x.shape[-1]
        w = hk.get_parameter(
            "w", [input_size, num_experts], jnp.float32, init=hk.initializers.Constant(0)
        )
        if sharding:
            w = with_sharding_constraint(w, sharding)

        out = jnp.dot(x, w.astype(fprop_dtype))
        return out


class MoELayer(hk.Module):
    def __init__(
        self,
        num_experts: int,
        layer_fn: Callable,
        router: Router,

View on GitHub (pinned to 7050ed204b)

Solutions

1. Inspect x.shape right before the MoE block in your forward pass; it must end with the hidden size (2048 for Grok-1), e.g. [batch, seq, 2048] or [batch, seq, wms, 2048].
2. Replace over-aggressive squeezes: use x = x[..., None, :] / indexing rather than x.squeeze(); restore dropped axes with x = x.reshape(1, 1, -1) or jnp.atleast_2d as appropriate.
3. Feed the model the same inputs run.py builds: a [batch, seq] integer token array produced by the SentencePiece tokenizer path, not a scalar token id.
4. Add a one-line assert in your wrapper: assert x.ndim >= 1 and x.shape[-1] == 2048 before calling the model.

Example fix

# before
token = jnp.array(prompt_ids[0])          # shape () after taking one element
out = model.apply(params, token)

# after
tokens = jnp.array(prompt_ids)[None, :]  # shape [1, T]
out = model.apply(params, tokens)

import jax.numpy as jnp

def check_router_input(x: jax.Array) -> None:
    # MoE router requires a trailing hidden dim (2048 for Grok-1)
    assert x.shape != (), f'scalar input to MoE router: {x.shape}'
    assert x.shape[-1] == 2048, f'unexpected hidden size: {x.shape[-1]}'
    x = jnp.atleast_2d(x)  # defensive normalization, shape [N, 2048]
    return x

def is_valid_moe_input(x: jax.Array) -> bool:
    """Router input must be non-scalar with a trailing feature axis."""
    return hasattr(x, 'shape') and len(x.shape) >= 1 and x.shape[-1] > 0

When it happens

Trigger: Calling the Grok-1 Transformer model (run.py sample or your own forward pass) where the tokens/embeddings entering a MixtureOfExperts block are 0-dimensional: typically x = x.squeeze() or jnp.squeeze applied too aggressively before the block, indexing with x[i] on a 1-d array, or feeding a batch of shape () because a collate/encoding step returned a scalar token id instead of a [batch, seq] array.

Common situations: Writing a custom inference wrapper that squeezes logits/embeddings per token; adapting the repo for batch size 1 and accidentally reducing [1, 1, 2048] all the way to a scalar; passing output of tokenizer id (a Python int converted with jnp.array without reshape) directly instead of the [B, T] token array the model expects.

Related errors

• Mask dimensionality {mask.ndim} must match logits dimensiona
• Parameters in the code are not matching checkpoint parameter

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