tracel-ai/burn · error

Can't differentiate avg pool 2d backward.

Error message

Can't differentiate avg pool 2d backward.

What it means

The autodiff backend does not implement the backward pass of avg_pool2d. Its module ops module implements only the forward pass; calling the backward entry point hits an unconditional panic because no gradient rule exists for this op on this backend combination.

Source

Thrown at crates/burn-autodiff/src/ops/module.rs:1498

                kernel_size,
                stride,
                padding,
                count_include_pad,
                ceil_mode,
            )),
        }
    }

    fn avg_pool2d_backward(
        _x: AutodiffTensor<B>,
        _grad: AutodiffTensor<B>,
        _kernel_size: [usize; 2],
        _stride: [usize; 2],
        _padding: [usize; 2],
        _count_include_pad: bool,
        _ceil_mode: bool,
    ) -> AutodiffTensor<B> {
        panic!("Can't differentiate avg pool 2d backward.");
    }

    fn max_pool1d(
        x: AutodiffTensor<B>,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        dilation: usize,
        ceil_mode: bool,
    ) -> AutodiffTensor<B> {
        match MaxPool1D
            .prepare::<C>([x.node.clone()])
            .compute_bound()
            .stateful()
        {
            OpsKind::Tracked(mut prep) => {
                let x_state = prep.checkpoint(&x);
                let settings = get_device_settings::<B>(&x.primitive.device());

View on GitHub (pinned to d16f7ba2ed)

Solutions

1. Use max_pool2d instead of avg_pool2d for pooling in layers that need gradients on this backend
2. Enable/upgrade the inner backend (e.g. burn-tch or burn-candle) so it supplies its own avg_pool2d backward
3. Downsample via strided conv2d or reshape+mean reduction composed of ops the autodiff backend supports
4. Implement avg_pool2d_backward for your backend by delegating to its primitives

Example fix

// before
let pooled = pool::avg_pool2d(&x, [2, 2], [2, 2], [0, 0], true, false);
// after
let pooled = pool::max_pool2d(&x, [2, 2], [2, 2], [0, 0], false);

// Before training, verify pooling layers use ops with backward support on autodiff
if uses_avg_pool2d(&model_config) {
    eprintln!("avg_pool2d backward is unimplemented in burn-autodiff; use max_pool2d or strided conv");
}

When it happens

Trigger: Calling burn_autodiff::module::avg_pool2d_backward directly, or a code path (e.g. a custom op or delegated backward) that requests the gradient of an avg_pool2d output from AutodiffTensor<B> where the inner backend does not provide it.

Common situations: Training a CNN with AvgPool2d layers on an autodiff-wrapped backend that lacks its own pool backward; upgrading burn to a version where the inner backend's pool backward is not yet implemented; hand-rolling backward passes against the module ops API.

Related errors

• Can't differentiate max pool2d with indices backward.
• Can't differentiate adaptive avg pool2d backward.
• Can't differentiate adaptive avg pool3d backward.
• Can't differentiate interpolate backward.
• unimplemented!("float_scatter with {other:?} update is not i

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