tracel-ai/burn · error

mean_dim: unsupported dtype {:?}

Error message

mean_dim: unsupported dtype {:?}

What it means

burn-flex's `mean_dim` computes a sum along a dimension then divides by the dim size; it supports float dtypes and integer dtypes (via widening sum + scalar divide). Unsupported dtypes like Bool or quantized hit the catch-all panic.

Source

Thrown at crates/burn-flex/src/ops/reduce.rs:344

            let mut tensor = sum_result;
            let data: &mut [u8] = tensor.storage_mut();
            for x in data.iter_mut() {
                *x = ((*x as u32) / divisor) as u8;
            }
            tensor
        }
        DType::U16 => {
            let divisor = dim_size as u32;
            let mut tensor = sum_result;
            let data: &mut [u16] = tensor.storage_mut();
            for x in data.iter_mut() {
                *x = ((*x as u32) / divisor) as u16;
            }
            tensor
        }
        DType::U32 => scalar_div::<u32>(sum_result, dim_size as u32),
        DType::U64 => scalar_div::<u64>(sum_result, dim_size as u64),
        _ => panic!("mean_dim: unsupported dtype {:?}", dtype),
    }
}

/// Product of all elements in a tensor, returning a scalar tensor.
pub fn prod(tensor: FlexTensor) -> FlexTensor {
    match tensor.dtype() {
        DType::F32 => prod_impl::<f32>(&tensor),
        DType::F64 => prod_impl::<f64>(&tensor),
        DType::F16 => reduce_scalar_half(&tensor, |a, b| a * b, 1.0, f16::to_f32, f16::from_f32),
        DType::BF16 => reduce_scalar_half(&tensor, |a, b| a * b, 1.0, bf16::to_f32, bf16::from_f32),
        DType::I8 => prod_impl_widening::<i8>(&tensor),
        DType::I16 => prod_impl_widening::<i16>(&tensor),
        DType::I32 => prod_impl_widening::<i32>(&tensor),
        DType::I64 => prod_impl::<i64>(&tensor),
        DType::U8 => prod_impl_widening::<u8>(&tensor),
        DType::U16 => prod_impl_widening::<u16>(&tensor),
        DType::U32 => prod_impl_widening::<u32>(&tensor),
        DType::U64 => prod_impl::<u64>(&tensor),

View on GitHub (pinned to d16f7ba2ed)

Solutions

1. Cast first: `mask.cast(DType::F32).mean_dim(1)` (float is usually what you want for means).
2. For integer means where truncation matters, sum as wider int then divide explicitly.
3. Dequantize quantized tensors before calling mean_dim.

Example fix

// before
let acc = mask.mean_dim(1); // Bool
// after
let acc = mask.cast(DType::F32).mean_dim(1);

assert!(matches!(t.dtype(), DType::F32 | DType::F64 | DType::F16 | DType::BF16 | DType::I8 | DType::I16 | DType::I32 | DType::I64 | DType::U8 | DType::U16 | DType::U32 | DType::U64), "mean_dim unsupported for {:?}", t.dtype());

fn is_mean_capable(d: DType) -> bool {
    !matches!(d, DType::Bool | DType::QFloat(_))
}

When it happens

Trigger: Calling `Tensor::mean_dim(dim)` on a Bool or quantized tensor, or any dtype outside the float/I/U integer set — e.g. averaging a boolean mask per row.

Common situations: Computing per-sample mean of boolean correctness masks without casting; mean over quantized activations; generic kernels that assume float inputs but receive bools.

Related errors

• sum: unsupported dtype {:?}
• sum_dim: unsupported dtype {:?}
• prod: unsupported dtype {:?}
• prod_dim: unsupported dtype {:?}
• max: unsupported dtype {:?}

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