Implement 3D pooling operations: avg_pool3d, max_pool3d, and adaptive_avg_pool3d

[MathijsdeBoer]

Pull Request Template

I'm opening this pull request early in response to #4213, though this is probably a little out of my depth so feel free to rip it apart.

Checklist

• Confirmed that cargo run-checks command has been executed.
• Made sure the book is up to date with changes in this PR.

Related Issues/PRs

Aims to implement #4213

Changes

Added 3D variants of the 3 available pooling operations, AvgPool3d, AdaptiveAvgPool3d and MaxPool3d. Also add a first go at a CubeCL implementation.

Testing

In general, the 2D unit tests were copied and adapted for 3D information.

LLM Disclosure

Claude Opus 4.5-preview was used to find appropriate locations for, and generate the main bulk of submitted code. I have compared and verified most of the code that was easily understood with the existing 2D myself. Though I am unfamiliar with some of the aspects of GPU programming/CubeCL.

[MathijsdeBoer]

As an aside, when going through the existing modules, I wondered if an API change would be beneficial regarding the dimensionality of some modules. For example, there is Conv1d, Conv2d and Conv3d. Would changing this to a single struct, Conv<D>, be a more ergonomic API? I ask because the Tensor struct already sort of does this with the <B, D> generics.

[MathijsdeBoer]

Ah, looks like there's been a few changes since my fork that have caused some changes to used functions, causing some compilation errors. Not sure when/if I'll have a chance to fix all of that in the upcoming Christmas time, unfortunately.

[laggui]

Ah, looks like there's been a few changes since my fork that have caused some changes to used functions, causing some compilation errors. Not sure when/if I'll have a chance to fix all of that in the upcoming Christmas time, unfortunately.

I didn't check the code, only fixed the compilation error locally for you since it was an easy fix 🙂

Will try to have a look at this before Christmas break myself, but just wanted to say thanks for including the LLM Disclosure !

LLM Disclosure

Claude Opus 4.5-preview was used to find appropriate locations for, and generate the main bulk of submitted code. I have compared and verified most of the code that was easily understood with the existing 2D myself. Though I am unfamiliar with some of the aspects of GPU programming/CubeCL.

Using coding tools isn't against the rules, and as a reviewer this already gives me helpful context. I think it's the first time I see a disclosure like this here, and I do appreciate it 🙏

[MathijsdeBoer]

I didn't check the code, only fixed the compilation error locally for you since it was an easy fix 🙂

Will try to have a look at this before Christmas break myself, but just wanted to say thanks for including the LLM Disclosure !

No worries about speed whatsoever! I have just started my own (sorely needed) Christmas break today!

Using coding tools isn't against the rules, and as a reviewer this already gives me helpful context. I think it's the first time I see a disclosure like this here, and I do appreciate it 🙏

I did it last time as well (#3407), and I prefer letting everyone know ahead of time. I'm sorely aware of the contentious reputation that LLMs have in programming, so I might as well not hide it.

I don't make it a habit to let LLMs write large chunks of my code, but my guess in this context was that, since a 3D version of an existing 1D/2D layer should just be a relatively minor iteration of the existing code, there shouldn't be too much opportunity for hallucinations/vibe-code-isms.

That said, if it does turn out to be below the standards of the project, I apologize for wasting time and would fully agree with a rejection pending a more human involved rewrite.

[codecov]

Codecov Report

❌ Patch coverage is 5.29954% with 1644 lines in your changes missing coverage. Please review.
✅ Project coverage is 68.22%. Comparing base (9ff428a) to head (30d1b4b).
⚠️ Report is 131 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/burn-autodiff/src/ops/module.rs	0.00%	222 Missing ⚠️
crates/burn-cubecl/src/kernel/pool/max_pool3d.rs	0.00%	203 Missing ⚠️
...burn-cubecl/src/kernel/pool/avg_pool3d_backward.rs	0.00%	155 Missing ⚠️
...burn-cubecl/src/kernel/pool/max_pool3d_backward.rs	0.00%	124 Missing ⚠️
crates/burn-tch/src/ops/module.rs	0.00%	122 Missing ⚠️
crates/burn-cubecl/src/kernel/pool/avg_pool3d.rs	0.00%	119 Missing ⚠️
...burn-cubecl/src/kernel/pool/adaptive_avg_pool3d.rs	0.00%	99 Missing ⚠️
...cl/src/kernel/pool/adaptive_avg_pool3d_backward.rs	0.00%	99 Missing ⚠️
crates/burn-cubecl/src/ops/module.rs	0.00%	92 Missing ⚠️
crates/burn-cubecl/src/kernel/pool/pool3d.rs	0.00%	78 Missing ⚠️
... and 9 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #4231      +/-   ##
==========================================
- Coverage   68.87%   68.22%   -0.65%     
==========================================
  Files        1405     1415      +10     
  Lines      167607   169369    +1762     
==========================================
+ Hits       115442   115560     +118     
- Misses      52165    53809    +1644     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[laggui]

Just had a quick look, the cubecl implementations look correct!

We need to add tests to validate the expected results for forward + backward ops with different configurations, similar to what is currently done for pool2d operations.

For reference:

• https://github.com/tracel-ai/burn/blob/main/crates/burn-backend-tests/tests/autodiff/avgpool2d.rs
• https://github.com/tracel-ai/burn/blob/main/crates/burn-backend-tests/tests/tensor/float/module/avgpool2d.rs

Expected values are typically computed using an existing framework (e.g. pytorch) as reference.

And once #4263 is completed & merged, we can add similar logic to support asymmetric padding.

[MathijsdeBoer]

Might take me a minute, I pulled in the most recent changes since my last update and now I have 151 errors to go through hahaha

Edit: Seems to be mostly usize <-> u32 conversions, I can manage that

[MathijsdeBoer]

Hey there, this PR has been languishing for a while, as I haven't had the energy to pick up this last bit while I am in the closing stages of my PhD. I apologize.

If someone else wants to pick up the slack in the meantime, feel free to. Otherwise, I'll get to it when I'm out of this period in my life.

[antimora]

Closing it as stale. Please open a new PR with the fresh rebase. If anyone else wants to pick, you may also submit a new PR