feat(cubecl): support AdaptiveAvgPool3d

[jcwal1516]

Pull Request Template

Checklist

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

cargo run-checks tests with Flex by default. Use cargo run-checks --backend <backend> when
targeting another backend. Consider running additional tests relevant to the crates changed.

No book change is required because AdaptiveAvgPool3d is already part of Burn's public tensor API; this PR adds CubeCL backend support for the existing operation.

Related Issues/PRs

• Closes [burn-cubecl] AdaptiveAvgPool3d #5288.
• Depends on tracel-ai/cubek#521 at revision 604d51ac0403ecbe06156c958ac458ca45a550e6.
• Cubek Update wgpu requirement from 0.16.1 to 0.17.0 #521 replaces the auto-closed tracel-ai/cubek#497.

Changes

• Dispatch CubeCL AdaptiveAvgPool3d forward and backward operations to the native Cubek kernels.
• Convert Burn's NCDHW tensors to the channels-last layout expected by Cubek and restore NCDHW output layout.
• Pass a channels-last input view for backward shape metadata and materialize the output gradient in the layout required by Cubek.
• Extend forward and backward backend tests with dynamic, asymmetric, multichannel, divisible, output-one, and larger-output cases.
• Pin Cubek to the exact implementation revision and regenerate the lockfile.

This does not change Burn's public API, fallback behavior for other backends, or quantization behavior.

Testing

Passed on Burn commit f62594eb975f04d2d95bb0c728bc1e47484dc588 against Cubek revision 604d51ac0403ecbe06156c958ac458ca45a550e6:

• macOS Metal, no fusion:
  BURN_DEVICE=metal cargo test --release -p burn-backend-tests --no-default-features --features metal,std adaptive_avgpool3d -- --nocapture
  — 32 passed
• macOS Metal with fusion:
  BURN_DEVICE=metal cargo test --release -p burn-backend-tests --no-default-features --features metal,std,fusion adaptive_avgpool3d -- --nocapture
  — 32 passed
• cargo run-checks --backend metal
• cargo xtask doc build
• cargo xtask doc tests
• cargo metadata --format-version 1 --locked --no-deps
• git diff origin/main...HEAD --check

cargo run-checks --backend metal passed on this exact Burn commit against the exact Cubek revision, including formatting, typos, audit, Clippy, no-std checks, and the complete Metal backend suites with and without fusion. Every AdaptiveAvgPool3d and FFT test passed.

Hosted Cubek CI on the exact pinned revision passed every job.

Hosted Burn CI on this exact commit passed every job: code quality; documentation build and tests; macOS Metal; stable and previous Linux backends, crates, examples, and no-std; stable Windows backends, crates, and examples; and coverage upload.

The hosted Linux coverage matrix does not exercise these backend-specific lines: its backend shard tests Ndarray and Flex, while its crates shard excludes burn-cubecl and burn-cubecl-fusion. The Metal suites above cover the changed CubeCL paths directly.

Before this rebase, the byte-identical AdaptiveAvgPool3d implementation and tests also passed targeted no-fusion and fusion runs on WSL2/Linux CUDA and native Windows CUDA — 32 passed in each configuration.

No benchmarks were added or used, and this PR makes no performance claim.

[codecov]

Codecov Report

❌ Patch coverage is 55.55556% with 60 lines in your changes missing coverage. Please review.
✅ Project coverage is 67.68%. Comparing base (13f0a12) to head (f62594e).

Files with missing lines	Patch %	Lines
crates/burn-cubecl/src/kernel/pool/base.rs	0.00%	55 Missing ⚠️
crates/burn-cubecl/src/ops/module.rs	0.00%	5 Missing ⚠️

❌ Your patch check has failed because the patch coverage (55.55%) 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 (67.68%) 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    #5372      +/-   ##
==========================================
- Coverage   67.71%   67.68%   -0.03%     
==========================================
  Files        1256     1256              
  Lines      194779   194878      +99     
==========================================
+ Hits       131888   131912      +24     
- Misses      62891    62966      +75     

☔ View full report in Codecov by Harness.
📢 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]

Burn changes LGTM, but this is blocked by the upstream PR in cubek. Will mark as a draft until that PR is re-opened & merged 👍