tracel-ai/burn · error

Can't differentiate adaptive avg pool2d backward.

Error message

Can't differentiate adaptive avg pool2d backward.

What it means

adaptive_avg_pool2d_backward in the autodiff backend panics unconditionally: there is no gradient rule for adaptive average pooling 2d on this backend. Only the forward adaptive_avg_pool2d is implemented.

Source

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

            .prepare::<C>([x.node.clone()])
            .compute_bound()
            .stateful()
        {
            OpsKind::Tracked(mut prep) => {
                let x_state = prep.checkpoint(&x);
                prep.finish(x_state, B::adaptive_avg_pool2d(x.primitive, output_size))
            }
            OpsKind::UnTracked(prep) => {
                prep.finish(B::adaptive_avg_pool2d(x.primitive, output_size))
            }
        }
    }

    fn adaptive_avg_pool2d_backward(
        _x: AutodiffTensor<B>,
        _grad: AutodiffTensor<B>,
    ) -> AutodiffTensor<B> {
        panic!("Can't differentiate adaptive avg pool2d backward.");
    }

    fn adaptive_avg_pool3d(x: AutodiffTensor<B>, output_size: [usize; 3]) -> AutodiffTensor<B> {
        #[derive(Debug)]
        struct AdaptiveAvgPool3D;

        impl<B: Backend> Backward<B, 1> for AdaptiveAvgPool3D {
            type State = NodeId;

            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
                let [node_parent] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);
                let state = checkpointer.retrieve_node_output(ops.state);

View on GitHub (pinned to d16f7ba2ed)

Solutions

1. Replace AdaptiveAvgPool2d with fixed-kernel avg_pool2d only if gradients are supported, or with a flatten+linear combination
2. Ensure the inner backend implements adaptive_avg_pool2d_backward and use it directly instead of autodiff
3. Compute the pooling as reshape/chunk/mean ops that autodiff can differentiate
4. Contribute/implement the backward via the existing backward-framework (Backward<B, D>)

Example fix

// before
let pooled = adaptive_avg_pool2d(&x, [7, 7]);
// after
let pooled = avg_pool2d(&x, [7, 7], [7, 7], [0, 0], true, false); // only if backward exists

// Guard adaptive pooling heads in trainable models
if model_uses_adaptive_avg_pool2d && is_training {
    eprintln!("adaptive_avg_pool2d backward panics in burn-autodiff; replace with fixed avg_pool2d or flatten+linear");
}

When it happens

Trigger: Training a model whose backward pass needs gradients through adaptive_avg_pool2d (e.g. Spatial Pyramid Pooling heads, variable-size input pooling) on the autodiff backend.

Common situations: SPP/pooled-output heads in vision models; dynamic image-size pipelines that use AdaptiveAvgPool2d before a classifier; porting torchvision-style models to burn.

Related errors

• Can't differentiate avg pool 2d backward.
• Can't differentiate max pool2d with indices 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/11e38209cb0c2402. Report an issue: GitHub.