tracel-ai/burn · error

max_pool2d_with_indices: unsupported dtype {:?}

Error message

max_pool2d_with_indices: unsupported dtype {:?}

What it means

max_pool2d_with_indices dispatches on the input tensor's dtype across F32/F64/F16/BF16 only. Any other dtype hits the catch-all match arm and panics before the op returns output and index tensors. This is the indices-returning variant of max pooling on the Flex backend.

Source

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

                ceil_mode,
            ),
            DType::F16 => pool::max_pool2d_with_indices_f16(
                x,
                kernel_size,
                stride,
                padding,
                dilation,
                ceil_mode,
            ),
            DType::BF16 => pool::max_pool2d_with_indices_bf16(
                x,
                kernel_size,
                stride,
                padding,
                dilation,
                ceil_mode,
            ),
            dtype => panic!("max_pool2d_with_indices: unsupported dtype {:?}", dtype),
        };
        if indices.dtype() != DType::from(indices_dtype) {
            indices = Flex::int_cast(indices, indices_dtype);
        }
        MaxPool2dWithIndices::new(output, indices)
    }

    fn max_pool2d_with_indices_backward(
        x: FloatTensor<Flex>,
        _kernel_size: [usize; 2],
        _stride: [usize; 2],
        _padding: [usize; 2],
        _dilation: [usize; 2],
        _ceil_mode: bool,
        output_grad: FloatTensor<Flex>,
        indices: IntTensor<Flex>,
    ) -> MaxPool2dBackward<Flex> {
        let x_grad = match x.dtype() {

View on GitHub (pinned to d16f7ba2ed)

Solutions

1. Cast the input to a supported float dtype, e.g. x.cast(DType::F32), before the call.
2. Fix the dtype at the tensor's origin so pooling always receives floats.
3. Add the missing dtype arm (pool::max_pool2d_with_indices_<dtype>) in crates/burn-flex/src/ops/module.rs.

Example fix

// before
let (out, idx) = x_int.max_pool2d_with_indices([2, 2], [2, 2], [0, 0], [1, 1]); // panics
// after
let (out, idx) = x_int.cast(DType::F32).max_pool2d_with_indices([2, 2], [2, 2], [0, 0], [1, 1]);

assert!(matches!(x.dtype(), burn::tensor::DType::F32 | burn::tensor::DType::F64 | burn::tensor::DType::F16 | burn::tensor::DType::BF16), "max_pool2d_with_indices 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)
}

// Cast before the indices-returning pool op:
let x = if is_float_dtype(x.dtype()) { x } else { x.cast(burn::tensor::DType::F32) };
let (out, idx) = x.max_pool2d_with_indices([2, 2], [2, 2], [0, 0], [1, 1]);

When it happens

Trigger: Calling Tensor::max_pool2d_with_indices on the Flex backend with a non-float (non F32/F64/F16/BF16) input tensor.

Common situations: Integer tensors from preprocessing entering max-pool-with-indices layers; dtype drift after model load or backend switch; mixing Int and Float tensors in a vision model.

Related errors

• avg_pool2d_backward: unsupported dtype {:?}
• adaptive_avg_pool2d: unsupported dtype {:?}
• adaptive_avg_pool2d_backward: unsupported dtype {:?}
• max_pool2d: 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/d9d668898b79274b. Report an issue: GitHub.