tracel-ai/burn · error

burn_flex::layer_norm: unsupported dtype {:?}

Error message

burn_flex::layer_norm: unsupported dtype {:?}

What it means

layer_norm in burn-flex supports F32, F64 natively and F16/BF16 via an upcast-to-f32 path (layer_norm_via_f32). Any other dtype reaches the catch-all panic arm. The library requires floating-point input because layer norm computes mean/variance with floating-point math and takes float gamma/beta parameters.

Source

Thrown at crates/burn-flex/src/ops/activation.rs:603

        assert!(
            beta_shape.len() == 1 && beta_shape[0] == d_model,
            "layer_norm: beta must be a 1-D tensor of length equal to last dim of input \
             (got shape {:?}, expected [{}])",
            beta_shape,
            d_model,
        );
    }

    match input.dtype() {
        DType::F32 => layer_norm_f32(input, gamma, beta, epsilon as f32),
        DType::F64 => layer_norm_f64(input, gamma, beta, epsilon),
        DType::F16 => {
            layer_norm_via_f32::<f16>(input, gamma, beta, epsilon, f16::to_f32, f16::from_f32)
        }
        DType::BF16 => {
            layer_norm_via_f32::<bf16>(input, gamma, beta, epsilon, bf16::to_f32, bf16::from_f32)
        }
        dtype => panic!("burn_flex::layer_norm: unsupported dtype {:?}", dtype),
    }
}

fn layer_norm_via_f32<E: burn_backend::Element + bytemuck::Pod + Copy>(
    input: FlexTensor,
    gamma: FlexTensor,
    beta: Option<FlexTensor>,
    epsilon: f64,
    to_f32: fn(E) -> f32,
    from_f32: fn(f32) -> E,
) -> FlexTensor {
    let input_f32 = crate::ops::module::cast_to_f32::<E>(input, to_f32);
    let gamma_f32 = crate::ops::module::cast_to_f32::<E>(gamma, to_f32);
    let beta_f32 = beta.map(|b| crate::ops::module::cast_to_f32::<E>(b, to_f32));
    let out = layer_norm_f32(input_f32, gamma_f32, beta_f32, epsilon as f32);
    crate::ops::module::cast_from_f32::<E>(out, from_f32)
}

View on GitHub (pinned to d16f7ba2ed)

Solutions

1. Cast the input (and gamma/beta) to a float dtype, e.g. input.cast(DType::F32), before calling layer_norm.
2. Audit the preceding op so activations stay in F16/BF16/F32/F64 through the normalization layer.
3. Wrap layer_norm in a helper that upcasts unsupported dtypes to F32 and downcasts after.

Example fix

// before
let out = layer_norm(x_i8, &gamma, &beta, 1e-5);
// after
let out = layer_norm(x_i8.cast(DType::F32), &gamma.cast(DType::F32), &beta.cast(DType::F32), 1e-5).cast(x_i8.dtype());

if !matches!(input.dtype(), DType::F32 | DType::F64 | DType::F16 | DType::BF16) {
    input = input.cast(DType::F32);
}
let out = layer_norm(input, &gamma, &beta, epsilon);

fn supports_layer_norm(d: DType) -> bool {
    matches!(d, DType::F32 | DType::F64 | DType::F16 | DType::BF16)
}

let out = std::panic::catch_unwind(|| layer_norm(input.clone(), &gamma, &beta, eps))
    .unwrap_or_else(|_| layer_norm(input.cast(DType::F32), &gamma, &beta, eps));

When it happens

Trigger: Calling burn_flex::ops::layer_norm (public) with input of an integer or bool dtype; also if gamma/beta and input dtypes are mismatched in a way that resolves to an unhandled dtype branch. Commonly after a quantize step that left activations in I8.

Common situations: Quantized transformer inference feeding int8 activations into LayerNorm; porting a model where the previous backend auto-cast; constructing normalization params as integers in tests.

Related errors

• softmax: unsupported dtype {:?}
• attention: unsupported dtype {:?}
• binary_op: unsupported dtype {:?}
• scalar_op: unsupported dtype {:?}
• int_binary_op: unsupported dtype {:?}

AI-assisted analysis of tracel-ai/burn@d16f7ba2ed (2026-09-05). Data as JSON: /api/errors/196622926b058dbd. Report an issue: GitHub.