Give flat-Array-backed ComponentArrays a Mooncake tangent type - #397
Draft
AstitvaAggarwal wants to merge 1 commit into
Draft
Give flat-Array-backed ComponentArrays a Mooncake tangent type#397AstitvaAggarwal wants to merge 1 commit into
ComponentArrays a Mooncake tangent type#397AstitvaAggarwal wants to merge 1 commit into
Conversation
Mirrors CuArray's tangent_type = self pattern instead of relying on Mooncake's generic derived struct, which downstream code (e.g. SciMLSensitivity's ForwardDiffSensitivity adjoint) has to specially unwrap. SubArray-backed ComponentArrays still use the existing generic derivation and rules.
|
this seems like it uses internals aggressively and likely should be moved to mooncake proper akin to the discussion in JuliaArrays/StaticArrays.jl#1343 ? |
ComponentArrays a Mooncake tangent type
AstitvaAggarwal
marked this pull request as draft
July 31, 2026 15:16
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Gives flat-Array-backed ComponentArrays their own Mooncake tangent type instead of relying on the generic auto-derived struct. Motivated by SciML/SciMLSensitivity.jl#1565, which needed a local workaround to unwrap ComponentVector's generic tangent shape. This fixes it at the source: a ComponentVector's cotangent is just a ComponentVector, so no unwrapping is needed for that code path.
Note: SubArray-backed ComponentArrays are untouched and still use the existing generic derivation.
Changes
tangent_type/fdata_type/rdata_type/tangent: the core declaration. ComponentArray is immutable, so these need to be declared explicitly.zero_tangent_internal,randn_tangent_internal,increment_internal!!, etc.lgetfield/_new_primitive rules, forward and reverse mode, similar to Mooncake's ownMemoryrule and CuArray's_new_rule.to_cr_tangent/increment_and_get_rdata!: the ChainRules interop bridge.pointer_from_objrefdoesn't work on an immutable type.Randomto[deps], needed byrandn_tangent_internal.Tested via
test_rulepasses, both AD modes, for identity, sum, a named-component function, and the lgetfield/new primitives directly.