tracel-ai/burn · error

adaptive_avg_pool2d_backward: unsupported dtype {:?}

Error message

adaptive_avg_pool2d_backward: unsupported dtype {:?}

What it means

adaptive_avg_pool2d_backward dispatches on the input tensor's dtype across the four supported float types (F32, F64, F16, BF16). If x has any other dtype, the catch-all match arm panics with this message. It guards the backward (gradient) path of adaptive average pooling.

Source

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

        match x.dtype() {
            DType::F32 => pool::adaptive_avg_pool2d_f32(x, output_size),
            DType::F64 => pool::adaptive_avg_pool2d_f64(x, output_size),
            DType::F16 => pool::adaptive_avg_pool2d_f16(x, output_size),
            DType::BF16 => pool::adaptive_avg_pool2d_bf16(x, output_size),
            dtype => panic!("adaptive_avg_pool2d: unsupported dtype {:?}", dtype),
        }
    }

    fn adaptive_avg_pool2d_backward(
        x: FloatTensor<Flex>,
        grad: FloatTensor<Flex>,
    ) -> FloatTensor<Flex> {
        match x.dtype() {
            DType::F32 => pool::adaptive_avg_pool2d_backward_f32(x, grad),
            DType::F64 => pool::adaptive_avg_pool2d_backward_f64(x, grad),
            DType::F16 => pool::adaptive_avg_pool2d_backward_f16(x, grad),
            DType::BF16 => pool::adaptive_avg_pool2d_backward_bf16(x, grad),
            dtype => panic!(
                "adaptive_avg_pool2d_backward: unsupported dtype {:?}",
                dtype
            ),
        }
    }

    fn adaptive_avg_pool3d(x: FloatTensor<Flex>, output_size: [usize; 3]) -> FloatTensor<Flex> {
        match x.dtype() {
            DType::F32 => pool::adaptive_avg_pool3d_f32(x, output_size),
            DType::F64 => pool::adaptive_avg_pool3d_f64(x, output_size),
            DType::F16 => pool::adaptive_avg_pool3d_f16(x, output_size),
            DType::BF16 => pool::adaptive_avg_pool3d_bf16(x, output_size),
            dtype => panic!("adaptive_avg_pool3d: unsupported dtype {:?}", dtype),
        }
    }

    fn adaptive_avg_pool3d_backward(
        x: FloatTensor<Flex>,

View on GitHub (pinned to d16f7ba2ed)

Solutions

  1. Cast x (and grad) to a supported float dtype, e.g. x.cast(DType::F32), before the backward pass.
  2. Fix the dtype at the source so the forward pass already receives float tensors.
  3. Add a match arm for the missing dtype in crates/burn-flex/src/ops/module.rs if support is required.

Example fix

// before
let grad_in = grad_int; // wrong dtype
let backward = adaptive_avg_pool2d_backward::<F32>(x, grad_in); // panics
// after
let backward = adaptive_avg_pool2d_backward::<F32>(x, grad_int.cast(DType::F32));
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), "adaptive_avg_pool2d_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

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

Prevention

When it happens

Trigger: Backpropagating through AdaptiveAvgPool2d on the Flex backend when the input tensor x (or its dtype slot) is not one of F32/F64/F16/BF16 — usually after an unintended integer cast upstream.

Common situations: Training a network where an integer tensor slipped into the pooling layer; checkpoint/precision mismatch; mixing Int tensors into a float compute graph.

Related errors


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