tracel-ai/burn · error

todo!("CubeCL backend does not yet support adaptive_avg_pool

Error message

todo!("CubeCL backend does not yet support adaptive_avg_pool3d_backward.")

What it means

Companion to 427: the CubeCL backend stubs `adaptive_avg_pool3d_backward` with `todo!`, so backpropagation through 3D adaptive average pooling panics. Even if a forward path were provided (e.g. another backend), training through this layer on CubeCL is impossible.

Source

Thrown at crates/burn-cubecl/src/ops/module.rs:303

        kernel::pool::adaptive_avg_pool2d(x, output_size)
    }

    fn adaptive_avg_pool2d_backward(
        x: FloatTensor<Self>,
        grad: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
        kernel::pool::adaptive_avg_pool2d_backward(x, grad)
    }

    fn adaptive_avg_pool3d(_x: FloatTensor<Self>, _output_size: [usize; 3]) -> FloatTensor<Self> {
        todo!("CubeCL backend does not yet support adaptive_avg_pool3d.")
    }

    fn adaptive_avg_pool3d_backward(
        _x: FloatTensor<Self>,
        _grad: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
        todo!("CubeCL backend does not yet support adaptive_avg_pool3d_backward.")
    }

    fn interpolate(
        x: FloatTensor<Self>,
        output_size: [usize; 2],
        options: InterpolateOptions,
    ) -> FloatTensor<Self> {
        kernel::interpolate::interpolate(x, output_size, options, Default::default()).unwrap()
    }

    fn interpolate_backward(
        x: FloatTensor<Self>,
        grad: FloatTensor<Self>,
        output_size: [usize; 2],
        options: InterpolateOptions,
    ) -> FloatTensor<Self> {
        kernel::interpolate::interpolate_backward(x, grad, output_size, options)
    }

View on GitHub (pinned to d16f7ba2ed)

Solutions

  1. Use fixed-kernel `avg_pool3d` (and its backward, which is implemented) instead of adaptive 3D pooling in trained models.
  2. Freeze/pool outside the graph: compute the pooled features on CPU (NdArray) and feed them as constants.
  3. Replace the layer with adaptive_avg_pool2d slices if the architecture allows.
  4. Track burn upstream for implementation of this backward op.

Example fix

// before
let pooled = x.adaptive_avg_pool3d([2, 2, 2]); // trained end-to-end, backward panics
// after
let pooled = x.avg_pool3d([4, 4, 4], [4, 4, 4], [0, 0, 0], false); // fixed pool with backward support
Defensive patterns

Strategy: fallback

Validate before calling

fn trainable_on_backend<B: Backend>() -> bool {
    // adaptive_avg_pool3d_backward unimplemented on CubeCL
    !std::any::type_name::<B>().contains("CubeCl")
}

Try / catch

// backward panics via todo!; avoid by not using AdaptiveAvgPool3d in trained CubeCL graphs
if !trainable_on_backend::<B>() {
    let pooled = x.avg_pool3d([4,4,4],[4,4,4],[0,0,0],false);
}

Prevention

When it happens

Trigger: Training (backward pass) a model containing an adaptive 3D average pooling layer on the CubeCL backend — autodiff calls `adaptive_avg_pool3d_backward` to compute input gradients.

Common situations: Fine-tuning 3D vision/video models with adaptive pooling heads on wgpu/CUDA; converting a PyTorch training loop with nn.AdaptiveAvgPool3d to burn on GPU backends.

Related errors


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