huggingface/candle · error

shape mismatch in rms-norm {:?} {:?}

Error message

shape mismatch in rms-norm {:?} {:?}

What it means

`candle_nn::ops::rms_norm` requires `alpha` to be a 1-D tensor whose length equals the last dimension (hidden size) of `xs`. Before dispatching to the kernel it checks `xs.dim(D::Minus1) == alpha.dims1()` and bails with both shapes when they differ. This is a shape-contract violation, not a memory/layout problem.

Source

Thrown at candle-nn/src/ops.rs:678

pub fn rms_norm_slow(x: &Tensor, alpha: &Tensor, eps: f32) -> Result<Tensor> {
    let x_dtype = x.dtype();
    let internal_dtype = match x_dtype {
        DType::F16 | DType::BF16 => DType::F32,
        d => d,
    };
    let hidden_size = x.dim(D::Minus1)?;
    let x = x.to_dtype(internal_dtype)?;
    let norm_x = (x.sqr()?.sum_keepdim(D::Minus1)? / hidden_size as f64)?;
    let x_normed = x.broadcast_div(&(norm_x + eps as f64)?.sqrt()?)?;
    x_normed.to_dtype(x_dtype)?.broadcast_mul(alpha)
}

pub fn rms_norm(xs: &Tensor, alpha: &Tensor, eps: f32) -> Result<Tensor> {
    let hidden_size_xs = xs.dim(D::Minus1)?;
    let hidden_size_alpha = alpha.dims1()?;
    if hidden_size_xs != hidden_size_alpha {
        candle::bail!(
            "shape mismatch in rms-norm {:?} {:?}",
            xs.shape(),
            alpha.shape()
        )
    }
    xs.apply_op2_no_bwd(alpha, &RmsNorm { eps })
}

#[derive(Debug, Clone)]
struct LayerNorm {
    eps: f32,
}

impl candle::CustomOp3 for LayerNorm {
    fn name(&self) -> &'static str {
        "layer-norm"
    }

View on GitHub (pinned to d5fee525bf)

Solutions

  1. Verify the alpha weight length matches the last dimension of the input; load the correct weight tensor for this layer.
  2. Reshape/transpose the input so its last dimension is the hidden size matching alpha.
  3. Ensure the model config's hidden_size matches the checkpoint being loaded.

Example fix

// before: alpha has 4096 elems, xs last dim is 3200
let out = rms_norm(&xs, &alpha_4096, 1e-6)?;
// after: use the weight matching the hidden size
let alpha = alpha_vars.get(("layers.0.input_layernorm", 3200))?;
let out = rms_norm(&xs, &alpha, 1e-6)?;
Defensive patterns

Strategy: validation

Validate before calling

// before calling rms_norm
let hidden = xs.dim(candle_core::D::Minus1)?;
let alpha_len = alpha.dims1()?; // also enforces 1-D alpha
if hidden != alpha_len {
    return Err(candle_core::Error::Msg(format!(
        "rms_norm: hidden size {hidden} != alpha len {alpha_len}"
    )));
}
let out = candle_nn::ops::rms_norm(&xs, &alpha, eps)?;

Type guard

fn alpha_matches(xs: &candle_core::Tensor, alpha: &candle_core::Tensor) -> candle_core::Result<bool> {
    Ok(alpha.dims().len() == 1 && alpha.dims()[0] == xs.dim(candle_core::D::Minus1)?)
}

Try / catch

match candle_nn::ops::rms_norm(&xs, &alpha, eps) {
    Ok(out) => out,
    Err(e) if e.to_string().contains("shape mismatch in rms-norm") => {
        // log both shapes, then fail fast with config context
        return Err(e);
    }
    Err(e) => return Err(e),
}

Prevention

When it happens

Trigger: Calling `rms_norm(xs, alpha, eps)` where `alpha`'s element count differs from `xs.dim(D::Minus1)` — e.g. weight from a model with hidden_size 4096 fed with activations of 3200, or alpha passed with 2-D shape (dims1() would also error) or transposed.

Common situations: Mixing layers/weights from different model configs, loading the wrong checkpoint, off-by-one reshaping so the last dim doesn't match the weight, or forgetting that rms_norm normalizes along the last dimension.

Understand the failure class

Background: Tensor shape mismatch errors ("must have shape", "expected shape ... got ..."): when tensor dimensions disagree with what an op or layer was told to expect — this error's family across 6 libraries.

Related errors


AI-assisted analysis of huggingface/candle@d5fee525bf (2026-09-02). Data as JSON: /api/errors/dcd9a41dd62301f7. Report an issue: GitHub.