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.32407% with 1636 lines in your changes missing coverage. Please review.
✅ Project coverage is 68.33%. Comparing base (91dd62c) to head (3addfaa).
⚠️ Report is 1 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%	122 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_ops.rs	0.00%	92 Missing ⚠️
crates/burn-cubecl/src/kernel/pool/pool3d.rs	0.00%	72 Missing ⚠️
... and 9 more

❌ Your patch check has failed because the patch coverage (5.32%) is below the target coverage (80.00%). You can increase the patch coverage or adjust the target coverage.
❌ Your project check has failed because the head coverage (68.33%) is below the target coverage (80.00%). You can increase the head coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #4231      +/-   ##
==========================================
- Coverage   68.96%   68.33%   -0.64%     
==========================================
  Files        1346     1360      +14     
  Lines      165578   167453    +1875     
==========================================
+ Hits       114192   114424     +232     
- Misses      51386    53029    +1643     

☔ 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.