tracel-ai/burn · error

adaptive_avg_pool2d: unsupported dtype {:?}

Error message

adaptive_avg_pool2d: unsupported dtype {:?}

What it means

adaptive_avg_pool2d in the burn-flex backend matches on the input tensor's dtype and only dispatches to the F32/F64/F16/BF16 kernel variants. Any other dtype hits the catch-all arm and panics. The backend deliberately fails fast instead of silently misinterpreting data.

Source

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

            DType::BF16 => pool::avg_pool2d_backward_bf16(
                x,
                grad,
                kernel_size,
                stride,
                padding,
                count_include_pad,
            ),
            dtype => panic!("avg_pool2d_backward: unsupported dtype {:?}", dtype),
        }
    }

    fn adaptive_avg_pool2d(x: FloatTensor<Flex>, output_size: [usize; 2]) -> FloatTensor<Flex> {
        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
            ),
        }
    }

View on GitHub (pinned to d16f7ba2ed)

Solutions

1. Cast the input to a supported float dtype before the call: x.cast(DType::F32).
2. Audit the pipeline producing the tensor so float dtype is guaranteed at the pooling layer.
3. Add the missing dtype arm (pool::adaptive_avg_pool2d_<dtype>) in crates/burn-flex/src/ops/module.rs if the backend should support it.

Example fix

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

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

// Panic-based; guard before the call:
let x = if is_float_dtype(x.dtype()) { x } else { x.cast(burn::tensor::DType::F32) };
let out = x.adaptive_avg_pool2d([1, 1]);

When it happens

Trigger: Calling adaptive_avg_pool2d (via Tensor::adaptive_avg_pool2d) on the Flex backend with a tensor whose dtype is not F32/F64/F16/BF16, e.g. an Int or Bool tensor.

Common situations: Feeding integer-encoded images or labels directly into an adaptive pooling layer; a preprocessing pipeline that forgot .float(); dtype drift after loading a model with a different precision config.

Related errors

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