Repository navigation
fix: materialize host views for AutoMooncake (CTBaseMooncake extension) - #632
Merged
Merged
Conversation
Mooncake returns a structural Mooncake.Tangent for a SubArray, even with friendly_tangents=true (chalk-lab/Mooncake.jl#1294), so the gradients of CTBase.Differentiation taken at a view of a host array could not be used as arrays (CTFlows.jl#448). - Differentiation: tags MooncakeTag / GenericADTag, _ad_tag(ad), and the _materialize(tag, x) hook (pass-through by default). - New CTBaseMooncake extension (weakdep Mooncake): _materialize copies a view of a host Array for MooncakeTag. - CTBaseDifferentiationInterface: every entry point adjusts its inputs (_ad_inputs) before building the closures, so the captured arguments are dense too (copying only the differentiated argument is 30-40% slower). Other backends and device arrays are unchanged. - Tests, docs wiring, changelog. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
ocots
added
run ci github-runner
Contributor
Breakage — packagesTested commit:
|
`using ..Core: AbstractTag` is not resolved by JET's top-level analysis on Julia 1.10; use `using ..Core: Core` and `Core.AbstractTag`, like the other submodules. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The `:ad_backend` option already holds an ADTypes backend, and ADTypes is a
hard dependency: CTBaseMooncake overloads
`_materialize(::ADTypes.AutoMooncake, ::SubArray{<:Any,<:Any,<:Array})`
directly. MooncakeTag, GenericADTag and _ad_tag are removed (ad_tags.jl ->
materialize.jl), and with them the import of Core.AbstractTag.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Closed
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.
Gradients of
CTBase.Differentiationtaken withAutoMooncakeat a view of a host array are now arrays. Mooncake returns a structuralMooncake.Tangent(fieldsparent,indices, …) for aSubArray, even withfriendly_tangents=true(chalk-lab/Mooncake.jl#1294), which madeFlow(h; ad_backend=AutoMooncake())fail on CPU in CTFlows (control-toolbox/CTFlows.jl#448, plan and probes there).What changes
src/Differentiation/materialize.jl(new): the hook_materialize(ad::ADTypes.AbstractADType, x), a pass-through by default. It dispatches on the ADTypes backend the:ad_backendoption already holds: no tag type is needed, ADTypes being a hard dependency.ext/CTBaseMooncake.jl(new extension, weak dependencyMooncake):_materialize(::ADTypes.AutoMooncake, x::SubArray{<:Any,<:Any,<:Array}) = copy(x).ext/CTBaseDifferentiationInterface.jl: every entry point (hamiltonian_gradient,variable_gradient, the threepseudo_*,gradient,differentiate,pushforward) adjusts its inputs with_ad_inputsbefore building the closures, so the captured arguments are dense too. Copying only the differentiated argument was measured 30–40 % slower.SubArrayof anArrayonly).test/suite/extensions/test_mooncake_extension.jl(57 tests: the hook, every entry point on vector and matrix views against ForwardDiff). Docs wiring, changelog.Checks run locally
test_extensions_quality(TST8).No breaking change. Next: release 0.30.13 (stable), then the compat bump in CTFlows.
🤖 Generated with Claude Code