What causes "Optional argument type mismatch" in tracel-ai/burn?

Convert the optional argument to the same dtype as the input, e.g. bias.cast:: () / create it with the same element type parameter.

tracel-ai/burn · error

Optional argument type mismatch

Error message

Optional argument type mismatch

What it means

The module_op! macro dispatches tensors by dtype (F32/F64) and maps an optional argument to the same dtype. If the optional argument tensor's dtype variant doesn't match the primary tensor's matched variant (e.g. bias f64 while input f32), the macro panics with 'Optional argument type mismatch'.

Source

Thrown at crates/burn-ndarray/src/ops/module.rs:39

};
use burn_backend::{
    TensorMetadata,
    ops::{attention::attention_fallback, conv::pad_asymmetric_conv_input, *},
    tensor::FloatTensor,
};
use burn_std::IntDType;

macro_rules! module_op {
    // Module op with inputs (inp), optional (opt) and arguments (args).
    // Converts NdArrayStorage to SharedArray for compatibility with existing operations.
    (inp($($x:tt),+), opt($($opt:tt),*), $element:ident, $op:expr) => {{
        #[allow(unused_parens, unreachable_patterns)]
        match ($($x),+) {
            ($(NdArrayTensor::F32($x)),+) => {
                type $element = f32;
                $op(
                    $($x.into_shared()),+
                    $(, $opt.map(|o| match o { NdArrayTensor::F32(val) => val.into_shared(), _ => panic!("Optional argument type mismatch") }))*
                )
            }
            ($(NdArrayTensor::F64($x)),+) => {
                type $element = f64;
                $op(
                    $($x.into_shared()),+
                    $(, $opt.map(|o| match o { NdArrayTensor::F64(val) => val.into_shared(), _ => panic!("Optional argument type mismatch") }))*
                )
            }
            _ => panic!("Data type mismatch"),
        }
    }};
}

impl ModuleOps<Self> for NdArray {
    fn conv2d(
        x: NdArrayTensor,
        weight: NdArrayTensor,

View on GitHub (pinned to d16f7ba2ed)

Solutions

Convert the optional argument to the same dtype as the input, e.g. bias.cast::<f32>() / create it with the same element type parameter.
Declare all tensors with the same element type: Tensor<NdArray<f32>, _> for both input and bias.
Check checkpoint/config for f64 weights and cast on load.
Make helper functions generic but instantiate with one E type.

Example fix

// before
let x: Tensor<NdArray<f32>, 3> = input;
let bias: Tensor<NdArray<f64>, 1> = Tensor::ones([c]);
conv2d(x, weight, Some(bias), options); // panic
// after
let bias: Tensor<NdArray<f32>, 1> = Tensor::ones([c]);
conv2d(x, weight, Some(bias), options);

Defensive patterns

Strategy: type-guard

Validate before calling

fn ensure_bias_dtype<E>(bias: &Tensor<NdArray<E>, 1>) { /* instantiate with same E as input */ }

Type guard

fn f64_to_f32(opt: Option<Tensor<NdArray<f64>, 1>>) -> Option<Tensor<NdArray<f32>, 1>> {
    opt.map(|b| b.cast::<f32>())
}

Prevention

Always construct bias/weights with the same element type as the model backend.
Avoid mixing f32 and f64 tensors in one model.
Cast checkpoint weights at load time to the model dtype.

When it happens

Trigger: Calling a module op (conv2d, conv_transpose2d, interpolate, etc.) where the main tensor(s) are F32 but the optional argument (like a bias tensor) is F64 (or vice versa).

Common situations: Creating bias with Tensor::full/ones defaulting to a different element type; loading weights from a checkpoint saved with f64; mixing Tensor<NdArray<f32>> and Tensor<NdArray<f64>> in one op call.

Related errors

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