rule for LuxLib.Impl.batchnorm_cudnn and fix CI Mooncake tests. - #1738
rule for LuxLib.Impl.batchnorm_cudnn and fix CI Mooncake tests.#1738AstitvaAggarwal wants to merge 3 commits into
LuxLib.Impl.batchnorm_cudnn and fix CI Mooncake tests.#1738Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request integrates Mooncake.jl into the LuxLib library, enabling reverse-mode automatic differentiation support for cuDNN-based batch normalization. The changes include adding Mooncake as a dependency, implementing a custom rrule for batchnorm_cudnn, and introducing a helper function _unwrap_mooncake_tangent in LuxTestUtils to handle tangent unwrapping. The review feedback correctly identifies a potential runtime error in the new test suite, suggesting the use of the newly implemented unwrapping utility to ensure compatibility with Mooncake's tangent objects.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| ) | ||
| ∂args_mooncake = ∂all_mooncake[2:end] # drop the gradient w.r.t. `f` itself | ||
| for (∂zyg, ∂mc) in zip(∂args_zygote, ∂args_mooncake) | ||
| @test Array(∂mc) ≈ Array(∂zyg) atol = atol rtol = rtol |
There was a problem hiding this comment.
The check_mooncake_correctness function directly calls Array(∂mc) on the result of Mooncake.value_and_gradient!!. If Mooncake returns a Tangent object instead of a raw array (which can happen depending on the configuration), Array(∂mc) will throw a MethodError. It is safer to use LuxTestUtils._unwrap_mooncake_tangent to ensure the tangent is unwrapped before converting to an Array.
@test Array(LuxTestUtils._unwrap_mooncake_tangent(∂mc)) ≈ Array(∂zyg) atol = atol rtol = rtol
667b04d to
6e301ec
Compare
1 participant
Adds a
Mooncakerrule!!forLuxLib.Impl.batchnorm_cudnn, enabling Mooncake AD through Lux'sGPUBatchNorm. The rule is caching the forward pass's batch stats so the pullback doesn't recompute them non-deterministically (similar pattern as FluxML/NNlib#727). New tests check both interface registration and correctness (vs Zygote).Note : Ive also fixed the LuxTestUtils bug that was causing Mooncake to fail completely on CI, due to which its testing was disabled temporarily (#1703 etc.)