tracel-ai/burn · error

Can't differentiate max pool2d with indices backward.

Error message

Can't differentiate max pool2d with indices backward.

What it means

max_pool2d_with_indices_backward in the autodiff backend is a stub that panics. The autodiff layer cannot differentiate max pooling with indices because the forward op's gradient rule relies on the inner backend's indices-based backward, which is unavailable here.

Source

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

                );
                let output_tensor = prep.finish(output.output);

                MaxPool2dWithIndices::new(output_tensor, output.indices)
            }
        }
    }

    fn max_pool2d_with_indices_backward(
        _x: AutodiffTensor<B>,
        _kernel_size: [usize; 2],
        _stride: [usize; 2],
        _padding: [usize; 2],
        _dilation: [usize; 2],
        _ceil_mode: bool,
        _output_grad: AutodiffTensor<B>,
        _indices: IntTensor<B>,
    ) -> MaxPool2dBackward<Self> {
        panic!("Can't differentiate max pool2d with indices backward.");
    }
    fn adaptive_avg_pool1d(x: AutodiffTensor<B>, output_size: usize) -> AutodiffTensor<B> {
        #[derive(Debug)]
        struct AdaptiveAvgPool1D;

        impl<B: Backend> Backward<B, 1> for AdaptiveAvgPool1D {
            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. Use max_pool2d (without indices) which has a working forward+backward in the autodiff backend
2. Call max_pool2d_with_indices_backward on the inner backend instead of the autodiff wrapper
3. Restructure the model to avoid differentiating through max pooling twice
4. Implement the backward by scattering the output grad into the input positions given by the indices tensor

Example fix

// before
let (out, idx) = max_pool2d_with_indices(&x, [2,2], [2,2], [0,0], [1,1], false);
// after
let out = max_pool2d(&x, [2,2], [2,2], [0,0], [1,1], false);

// Avoid indices-based pooling in differentiable graphs
if config.pooling == Pooling::MaxWithIndices && config.training {
    eprintln!("max_pool2d_with_indices_backward panics under autodiff; use max_pool2d");
}

When it happens

Trigger: Calling max_pool2d_with_indices_backward on an AutodiffTensor, or hitting a backward pass that routes through this stub (e.g. second-order differentiation of max_pool2d).

Common situations: Differentiating through a network containing MaxPool2d twice (double-backward); building custom autodiff ops that reuse the pooled indices; tests that call the backward API directly.

Related errors

• Can't differentiate avg pool 2d 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/b06c3020e390757e. Report an issue: GitHub.