max_pool2d_with_indices_backward: unsupported dtype {:?}

Error message

max_pool2d_with_indices_backward: unsupported dtype {:?}

What it means

max_pool2d_with_indices_backward computes the input gradient by dispatching on x.dtype(), supporting only F32, F64, F16 and BF16. Any other dtype reaches the catch-all arm and panics. This guards the backward pass of max pooling with indices on the Flex backend.

Source

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

        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() {
            DType::F32 => pool::max_pool2d_backward_f32(x, output_grad, indices),
            DType::F64 => pool::max_pool2d_backward_f64(x, output_grad, indices),
            DType::F16 => pool::max_pool2d_backward_f16(x, output_grad, indices),
            DType::BF16 => pool::max_pool2d_backward_bf16(x, output_grad, indices),
            dtype => panic!(
                "max_pool2d_with_indices_backward: unsupported dtype {:?}",
                dtype
            ),
        };
        MaxPool2dBackward::new(x_grad)
    }

    fn interpolate(
        x: FloatTensor<Flex>,
        output_size: [usize; 2],
        options: InterpolateOptions,
    ) -> FloatTensor<Flex> {
        match (options.mode, x.dtype()) {
            (InterpolateMode::Nearest, DType::F32) => {
                interpolate::interpolate_nearest_f32(x, output_size, options.align_corners)
            }
            (InterpolateMode::Nearest, DType::F64) => {
                interpolate::interpolate_nearest_f64(x, output_size, options.align_corners)

View on GitHub (pinned to d16f7ba2ed)

Solutions

Cast x/output_grad/indices appropriately so x is a supported float dtype before backward.
Ensure the forward pass only receives float tensors so the saved state is float.
Add a match arm calling pool::max_pool2d_backward_<dtype> in crates/burn-flex/src/ops/module.rs if a new dtype is needed.

Example fix

// before
let dx = max_pool2d_with_indices_backward::<F32>(x_int, grad, idx); // panics
// after
let dx = max_pool2d_with_indices_backward::<F32>(x_int.cast(DType::F32), grad, idx);

Defensive patterns

Strategy: validation

Validate before calling

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

Type guard

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)
}

Try / catch

// Ensure float input before backward:
let x = if is_float_dtype(x.dtype()) { x } else { x.cast(burn::tensor::DType::F32) };

Prevention

Confirm saved max-pool inputs are float before calling backward.
Match precision settings between training config and checkpoint.
Run a full forward/backward smoke test after backend or dtype changes.

When it happens

Trigger: Backpropagating through MaxPool2d (with indices) on the Flex backend when the saved input x has a dtype other than F32/F64/F16/BF16.

Common situations: Training where an integer tensor entered max pooling; precision mismatch between forward-saved tensors and backward expectations; pipeline bugs feeding Int data into vision models.

Related errors

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