What causes "shape mismatch in layer-norm src: {:?} alpha: {:?} beta: {:?}" in huggingface/candle?

Verify the last dim of xs equals alpha.len() and beta.len() before the call

shape mismatch in layer-norm src: {:?} alpha: {:?} beta: {:?

Error message

shape mismatch in layer-norm src: {:?} alpha: {:?} beta: {:?}

What it means

layer_norm validates that the last dimension of the input tensor equals the length (dims1) of both the alpha (scale) and beta (bias) tensors. The elementwise affine parameters must span exactly the normalized (hidden) dimension; if not, the shapes are printed and an error is raised before the op is dispatched.

Source

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

    let x = x.to_dtype(internal_dtype)?;
    let x = {
        let mean_x = (x.sum_keepdim(D::Minus1)? / hidden_size as f64)?;
        x.broadcast_sub(&mean_x)?
    };
    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)?
        .broadcast_add(beta)
}

pub fn layer_norm(xs: &Tensor, alpha: &Tensor, beta: &Tensor, eps: f32) -> Result<Tensor> {
    let hidden_size_xs = xs.dim(D::Minus1)?;
    let hidden_size_alpha = alpha.dims1()?;
    let hidden_size_beta = beta.dims1()?;
    if hidden_size_xs != hidden_size_alpha || hidden_size_xs != hidden_size_beta {
        candle::bail!(
            "shape mismatch in layer-norm src: {:?} alpha: {:?} beta: {:?}",
            xs.shape(),
            alpha.shape(),
            beta.shape()
        )
    }
    xs.apply_op3_no_bwd(alpha, beta, &LayerNorm { eps })
}

// https://pytorch.org/docs/stable/generated/torch.nn.PixelShuffle.html
pub fn pixel_shuffle(xs: &Tensor, upscale_factor: usize) -> Result<Tensor> {
    let (b_size, c, h, w) = xs.dims4()?;
    let out_c = c / upscale_factor / upscale_factor;
    xs.reshape((b_size, out_c, upscale_factor, upscale_factor, h, w))?
        .permute((0, 1, 4, 2, 5, 3))?
        .reshape((b_size, out_c, h * upscale_factor, w * upscale_factor))
}

View on GitHub (pinned to d5fee525bf)

Solutions

Verify the last dim of xs equals alpha.len() and beta.len() before the call
Ensure alpha and beta are 1-D tensors of size hidden_size (use dims1-compatible tensors)
Check the checkpoint/config hidden_size matches the model that produced xs
Use Tensor::broadcast-agnostic reshape: keep weight as shape (hidden_size,) not (1, hidden_size)

Example fix

// before
let alpha = Tensor::new(vec![0.0f32; 768], &dev)?; // wrong size
let out = layer_norm(&xs, &alpha, &beta, 1e-5)?;
// after
let hidden = xs.dim(candle::D::Minus1)?;
let alpha = Tensor::new(vec![0.0f32; hidden], &dev)?;
let out = layer_norm(&xs, &alpha, &beta, 1e-5)?;

Defensive patterns

Strategy: validation

Validate before calling

fn check_layernorm_shapes(xs: &Tensor, alpha: &Tensor, beta: &Tensor) -> candle::Result<()> {
    let h = xs.dim(candle::D::Minus1)?;
    if alpha.dims1()? != h || beta.dims1()? != h {
        candle::bail!("layer_norm shape mismatch: xs hidden={h} alpha={:?} beta={:?}", alpha.shape(), beta.shape());
    }
    Ok(())
}

Type guard

fn layernorm_shapes_ok(xs: &Tensor, alpha: &Tensor, beta: &Tensor) -> bool {
    xs.dim(candle::D::Minus1).map(|h| alpha.dims1() == Ok(h) && beta.dims1() == Ok(h)).unwrap_or(false)
}

Try / catch

match layer_norm(&xs, &alpha, &beta, eps) {
    Ok(y) => y,
    Err(e) if e.to_string().contains("shape mismatch in layer-norm") => {
        // log alpha/beta shapes and the expected hidden size, then re-raise with context
        Err(e)
    }
    Err(e) => Err(e),
}

Prevention

Keep alpha/beta as 1-D tensors of exactly hidden_size
Cross-check hidden_size between model config and checkpoint weights
Don't pass 2-D or broadcast-shaped weight matrices to layer_norm
Unit-test norm layers with the real checkpoint shapes at startup

When it happens

Trigger: Calling candle_nn::ops::layer_norm(xs, alpha, beta, eps) where xs.dim(D::Minus1) != alpha.dims1() or != beta.dims1(), e.g. weight tensors from a different hidden size than the activation tensor.

Common situations: Loading weights from a model checkpoint whose hidden_size differs from the model config; passing full 2D weight matrices instead of 1D vectors; mixing up layers (e.g. MLP weights fed into a norm layer); reshaping errors that change the last dim.

Related errors

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