tracel-ai/burn · error

Can't differentiate adaptive avg pool3d backward.

Error message

Can't differentiate adaptive avg pool3d backward.

What it means

adaptive_avg_pool3d_backward is a stub that panics: the autodiff backend provides no gradient rule for adaptive average pooling in 3D. The forward op works, but any training step that backpropagates through it fails at runtime.

Source

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

            .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_pool3d(x.primitive, output_size))
            }
            OpsKind::UnTracked(prep) => {
                prep.finish(B::adaptive_avg_pool3d(x.primitive, output_size))
            }
        }
    }

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

    fn interpolate(
        x: AutodiffTensor<B>,
        output_size: [usize; 2],
        options: InterpolateOptions,
    ) -> AutodiffTensor<B> {
        #[derive(Debug)]
        struct Interpolate;
        impl<B: Backend> Backward<B, 1> for Interpolate {
            type State = (NodeId, [usize; 2], InterpolateOptions);

            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {

View on GitHub (pinned to d16f7ba2ed)

Solutions

1. Use a fixed-kernel avg_pool3d whose gradient is implemented, if input sizes allow
2. Reshape/volume-mean decomposition using supported ops (slice/reshape/mean) so autodiff can differentiate
3. Run the pooling forward on the inner backend and stop gradients at that boundary
4. Implement adaptive_avg_pool3d_backward in the backend using the burn-autodiff Backward framework

Example fix

// before
let pooled = adaptive_avg_pool3d(&x, [1, 4, 4]);
// after
let pooled = avg_pool3d(&x, [2, 8, 8], [2, 8, 8], [0, 0, 0], true, false); // sizes known at compile time

if model_uses_adaptive_avg_pool3d && is_training {
    eprintln!("adaptive_avg_pool3d backward panics in burn-autodiff; use fixed avg_pool3d");
}

When it happens

Trigger: Backpropagating through adaptive_avg_pool3d (e.g. video or volumetric model heads pooling to a fixed output size) using the autodiff backend.

Common situations: 3D CNNs (video classification, medical volume models) with adaptive pooling heads; converting PyTorch models using AdaptiveAvgPool3d to burn.

Related errors

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