tracel-ai/burn · error

avg_pool2d_backward: unsupported dtype {:?}

Error message

avg_pool2d_backward: unsupported dtype {:?}

What it means

The burn-flex backend's avg_pool2d_backward dispatches on the gradient tensor's dtype and only implements F32, F64, F16 and BF16. If the tensor carries any other dtype (e.g. an integer or bool dtype that ended up in a float slot), the catch-all match arm panics with this message. It is an exhaustive-dispatch guard, not a recoverable error.

Source

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

                count_include_pad,
            ),
            DType::F16 => pool::avg_pool2d_backward_f16(
                x,
                grad,
                kernel_size,
                stride,
                padding,
                count_include_pad,
            ),
            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() {

View on GitHub (pinned to d16f7ba2ed)

Solutions

1. Check tensor.dtype() before the pooling backward call and cast to a supported float dtype with tensor.cast(DType::F32) (or .float()).
2. Fix the source of the wrong dtype upstream (e.g. an int_cast, integer input pipeline, or loaded checkpoint) so float tensors reach AvgPool2d.
3. If you need another float dtype implemented, add a match arm calling pool::avg_pool2d_backward_<dtype> in crates/burn-flex/src/ops/module.rs.

Example fix

// before
let grad = grad_int; // DType::I32
let x_grad = avg_pool2d_backward(x, grad, ...); // panics
// after
let grad = grad_int.cast(DType::F32);
let x_grad = avg_pool2d_backward(x, grad, ...);

assert!(matches!(grad.dtype(), burn::tensor::DType::F32 | burn::tensor::DType::F64 | burn::tensor::DType::F16 | burn::tensor::DType::BF16), "avg_pool2d_backward needs a float tensor, got {:?}", grad.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)
}

// Panics are not catchable in Rust; validate dtype first:
let grad = if is_float_dtype(grad.dtype()) { grad } else { grad.cast(burn::tensor::DType::F32) };

When it happens

Trigger: Calling the Burn Tensor API's avg_pool2d backward (e.g. training through AvgPool2d) on the Flex backend with a tensor whose dtype is not one of F32/F64/F16/BF16 — typically after an int_cast or a model that produced integer tensors feeding the pooling layer.

Common situations: Integer/quantized tensors passed where floats are expected; a dtype mismatch after loading weights or checkpoints saved with a different dtype; mixing backends where an Int tensor is accidentally fed into a float op.

Related errors

• adaptive_avg_pool2d: 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/5bc4c6320a644d90. Report an issue: GitHub.