What causes "avg_pool2d: unsupported dtype {:?}" in tracel-ai/burn?

Cast the input to a float dtype before avg_pool2d: x.cast(DType::F32).

avg_pool2d: unsupported dtype {:?}

Error message

avg_pool2d: unsupported dtype {:?}

What it means

avg_pool2d in the burn-flex module ops matches the input dtype (F32/F64/F16/BF16 supported) and delegates to the corresponding typed pooling kernel; any other dtype panics. Pooling ops operate on float activations, so an integer or bool tensor reaching avg_pool2d is a dtype-flow bug in the calling model.

Source

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

                ceil_mode,
            ),
            DType::F16 => pool::avg_pool2d_f16(
                x,
                kernel_size,
                stride,
                padding,
                count_include_pad,
                ceil_mode,
            ),
            DType::BF16 => pool::avg_pool2d_bf16(
                x,
                kernel_size,
                stride,
                padding,
                count_include_pad,
                ceil_mode,
            ),
            dtype => panic!("avg_pool2d: unsupported dtype {:?}", dtype),
        }
    }

    fn avg_pool2d_backward(
        x: FloatTensor<Flex>,
        grad: FloatTensor<Flex>,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
        count_include_pad: bool,
        _divisor_override: bool,
    ) -> FloatTensor<Flex> {
        match x.dtype() {
            DType::F32 => pool::avg_pool2d_backward_f32(
                x,
                grad,
                kernel_size,
                stride,

View on GitHub (pinned to d16f7ba2ed)

Solutions

Cast the input to a float dtype before avg_pool2d: x.cast(DType::F32).
Check .dtype() on the tensor entering the pool to identify which upstream op changed it to a non-float type.
Move any quantization/dequantization boundary so activations stay float through the pooling stage.
Add a new dtype arm in crates/burn-flex/src/ops/module.rs avg_pool2d if support is genuinely needed.

Example fix

// before
let pooled = avg_pool2d(features_i8, [2, 2], [2, 2], [0, 0], true, false);
// panic: avg_pool2d: unsupported dtype I8

// after
let x = features_i8.cast(burn::tensor::DType::F32);
let pooled = avg_pool2d(x, [2, 2], [2, 2], [0, 0], true, false);

Defensive patterns

Strategy: validation

Validate before calling

let x = if matches!(x.dtype(), DType::F32 | DType::F64 | DType::F16 | DType::BF16) { x } else { x.cast(DType::F32) };
let pooled = avg_pool2d(x, kernel_size, stride, padding, count_include_pad, ceil_mode);

Type guard

fn is_float(t: &Tensor<Flex>) -> bool {
    matches!(t.dtype(), DType::F32 | DType::F64 | DType::F16 | DType::BF16)
}

Try / catch

let pooled = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| avg_pool2d(x.clone(), ks, st, pd, cip, cm)))
    .unwrap_or_else(|_| avg_pool2d(x.cast(DType::F32), ks, st, pd, cip, cm));

Prevention

Keep activations float through pooling stages; place quantize boundaries elsewhere.
Find any round/clip-to-int helpers inserted before pooling and remove or move them.
Assert backbone output dtype before the classification head's pooling.
Check exported graph dtype changes around pool nodes after conversions.

When it happens

Trigger: Calling avg_pool2d (or an AvgPool2d module forward) with a non-float input tensor; pooling over quantized int8 feature maps without dequantization; passing integer label/mask tensors through a pooling stage.

Common situations: Classification heads pooling int-quantized backbone outputs; pipelines where a cast to integer (e.g. round/clip helper) was inserted before the pool; exported graphs with dtype changes around pooling.

Related errors

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