{"record":{"id":"32f32deb733c7901","repo":"xai-org/grok-1","slug":"input-must-not-be-scalar","errorCode":null,"errorMessage":"Input must not be scalar.","messagePattern":"Input must not be scalar\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"model.py","lineNumber":259,"sourceCode":"        routing_logits = self._router_weights(inputs, num_experts, sharding=P(\"data\"))\n        assert routing_logits.dtype == jnp.float32\n        routing_probs = jax.nn.softmax(routing_logits)\n\n        if padding_mask is not None:\n            routing_probs *= padding_mask\n\n        return routing_probs, routing_logits, 0\n\n    @hk.transparent\n    def _router_weights(\n        self,\n        x: jax.Array,\n        num_experts: int,\n        sharding: Optional[P] = None,\n    ):\n        fprop_dtype = x.dtype\n        if not x.shape:\n            raise ValueError(\"Input must not be scalar.\")\n\n        input_size = self.input_size = x.shape[-1]\n        w = hk.get_parameter(\n            \"w\", [input_size, num_experts], jnp.float32, init=hk.initializers.Constant(0)\n        )\n        if sharding:\n            w = with_sharding_constraint(w, sharding)\n\n        out = jnp.dot(x, w.astype(fprop_dtype))\n        return out\n\n\nclass MoELayer(hk.Module):\n    def __init__(\n        self,\n        num_experts: int,\n        layer_fn: Callable,\n        router: Router,","sourceCodeStart":241,"sourceCodeEnd":277,"githubUrl":"https://github.com/xai-org/grok-1/blob/7050ed204b8206bb8645c7b7bbef7252f79561b0/model.py#L241-L277","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["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].","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.","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.","Add a one-line assert in your wrapper: assert x.ndim >= 1 and x.shape[-1] == 2048 before calling the model."],"exampleFix":"# before\ntoken = jnp.array(prompt_ids[0])          # shape () after taking one element\nout = model.apply(params, token)\n\n# after\ntokens = jnp.array(prompt_ids)[None, :]  # shape [1, T]\nout = model.apply(params, tokens)","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\n\ndef check_router_input(x: jax.Array) -> None:\n    # MoE router requires a trailing hidden dim (2048 for Grok-1)\n    assert x.shape != (), f'scalar input to MoE router: {x.shape}'\n    assert x.shape[-1] == 2048, f'unexpected hidden size: {x.shape[-1]}'\n    x = jnp.atleast_2d(x)  # defensive normalization, shape [N, 2048]\n    return x","typeGuard":"def is_valid_moe_input(x: jax.Array) -> bool:\n    \"\"\"Router input must be non-scalar with a trailing feature axis.\"\"\"\n    return hasattr(x, 'shape') and len(x.shape) >= 1 and x.shape[-1] > 0","tryCatchPattern":null,"preventionTips":["Never use bare .squeeze() on activations flowing into the model; always pass an axis.","Keep token arrays rank-2 ([batch, seq]) end to end in generation loops; slice logits with [:, -1, :] instead of scalarizing.","Assert x.ndim >= 1 and x.shape[-1] == 2048 at the boundary of your inference wrapper."],"tags":["grok-1","jax","moe","router","input-validation","shape-error"],"backgroundTag":null,"analyzedSha":"7050ed204b8206bb8645c7b7bbef7252f79561b0","analyzedAt":"2026-08-15T04:18:25.087Z","schemaVersion":2},"datasetVersion":"2026-08-15T22:17:37.221Z"}