tracel-ai/burn · error

max_pool2d: unsupported dtype {:?}

Error message

max_pool2d: unsupported dtype {:?}

What it means

max_pool2d in the burn-flex backend dispatches on the input dtype and implements only F32, F64, F16 and BF16 variants. Any other dtype reaches the catch-all arm and panics. It is a fail-fast guard ensuring pooling kernels never receive non-float data.

Source

Thrown at crates/burn-flex/src/ops/module.rs:472

        stride: [usize; 2],
        padding: [usize; 2],
        dilation: [usize; 2],
        ceil_mode: bool,
    ) -> FloatTensor<Flex> {
        match x.dtype() {
            DType::F32 => {
                pool::max_pool2d_f32(x, kernel_size, stride, padding, dilation, ceil_mode)
            }
            DType::F64 => {
                pool::max_pool2d_f64(x, kernel_size, stride, padding, dilation, ceil_mode)
            }
            DType::F16 => {
                pool::max_pool2d_f16(x, kernel_size, stride, padding, dilation, ceil_mode)
            }
            DType::BF16 => {
                pool::max_pool2d_bf16(x, kernel_size, stride, padding, dilation, ceil_mode)
            }
            dtype => panic!("max_pool2d: unsupported dtype {:?}", dtype),
        }
    }

    fn max_pool2d_with_indices(
        x: FloatTensor<Flex>,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
        dilation: [usize; 2],
        ceil_mode: bool,
        indices_dtype: IntDType,
    ) -> MaxPool2dWithIndices<Flex> {
        let (output, mut indices) = match x.dtype() {
            DType::F32 => pool::max_pool2d_with_indices_f32(
                x,
                kernel_size,
                stride,
                padding,

View on GitHub (pinned to d16f7ba2ed)

Solutions

1. Cast the input to a supported float dtype before pooling: x.cast(DType::F32).
2. Fix the upstream producer of the tensor so float dtype is guaranteed.
3. Add a match arm for the needed dtype calling pool::max_pool2d_<dtype> in crates/burn-flex/src/ops/module.rs.

Example fix

// before
let pooled = feats_i32.max_pool2d([2, 2], [2, 2], [0, 0], [1, 1], false); // panics
// after
let pooled = feats_i32.cast(DType::F32).max_pool2d([2, 2], [2, 2], [0, 0], [1, 1], false);

assert!(matches!(x.dtype(), burn::tensor::DType::F32 | burn::tensor::DType::F64 | burn::tensor::DType::F16 | burn::tensor::DType::BF16), "max_pool2d needs a float tensor, got {:?}", x.dtype());

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

// Normalize dtype before pooling:
let x = if is_float_dtype(x.dtype()) { x } else { x.cast(burn::tensor::DType::F32) };
let pooled = x.max_pool2d([2, 2], [2, 2], [0, 0], [1, 1], false);

When it happens

Trigger: Calling Tensor::max_pool2d on the Flex backend with a tensor whose dtype is not one of F32/F64/F16/BF16, e.g. an Int tensor.

Common situations: Feeding integer feature maps (e.g. from a quantized or index-based stage) into MaxPool2d; forgetting .float() after an image pipeline; dtype changes when switching backends or loading checkpoints.

Related errors

• avg_pool2d_backward: unsupported dtype {:?}
• adaptive_avg_pool2d: unsupported dtype {:?}
• adaptive_avg_pool2d_backward: unsupported dtype {:?}
• max_pool2d_with_indices: unsupported dtype {:?}
• max_pool2d_with_indices_backward: unsupported dtype {:?}

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