Skip to content

Give flat-Array-backed ComponentArrays a Mooncake tangent type - #397

Draft
AstitvaAggarwal wants to merge 1 commit into
SciML:mainfrom
AstitvaAggarwal:mooncake-native-tangent-type
Draft

Give flat-Array-backed ComponentArrays a Mooncake tangent type#397
AstitvaAggarwal wants to merge 1 commit into
SciML:mainfrom
AstitvaAggarwal:mooncake-native-tangent-type

Conversation

@AstitvaAggarwal

@AstitvaAggarwal AstitvaAggarwal commented Jul 31, 2026

Copy link
Copy Markdown
Member

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.
  • The standard tangent operations Mooncake requires: zero_tangent_internal, randn_tangent_internal, increment_internal!!, etc.
  • lgetfield/_new_ primitive rules, forward and reverse mode, similar to Mooncake's own Memory rule and CuArray's _new_ rule.
  • to_cr_tangent/increment_and_get_rdata!: the ChainRules interop bridge.
  • Test-infrastructure overrides for Mooncake's TestUtils, adapted because pointer_from_objref doesn't work on an immutable type.
  • Added Random to [deps], needed by randn_tangent_internal.

Tested via

  • test_rule passes, both AD modes, for identity, sum, a named-component function, and the lgetfield/new primitives directly.
  • The SciMLSensitivity #1548 ODE repro produces a correct gradient end-to-end.
  • setindex!, reshape, copy, similar, vcat, map, two-CV arithmetic, a real ComponentMatrix, and Complex element types all checked by hand.
  • Nested/SubArray-backed access confirmed unaffected.

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.
@wsmoses

wsmoses commented Jul 31, 2026

Copy link
Copy Markdown

this seems like it uses internals aggressively and likely should be moved to mooncake proper akin to the discussion in JuliaArrays/StaticArrays.jl#1343 ?

@AstitvaAggarwal AstitvaAggarwal changed the title Give flat-Array-backed ComponentArrays their own Mooncake tangent type Give flat-Array-backed ComponentArrays a Mooncake tangent type Jul 31, 2026
@AstitvaAggarwal
AstitvaAggarwal marked this pull request as draft July 31, 2026 15:16
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants