Repository navigation
Add Mooncake extension marking searches inactive - #117
Merged
Merged
Conversation
Mirrors the existing EnzymeCore extension: searchsorted_last, searchsorted_first, searchsortedlast!, searchsortedfirst!, searchsortedrange, findequal, findfirstequal, and findfirstsortedequal all return a discrete index with no meaningful derivative, so mark them @zero_derivative for Mooncake too. This also works around a real Mooncake limitation: KIND_SIMD_LINEAR_SCAN's kernel embeds a raw llvmcall SIMD intrinsic that Mooncake's intrinsic-lifting machinery cannot translate. Marking the dispatcher itself inactive stops Mooncake before it ever descends into that (or any other strategy's) kernel. Surfaced via native Hessian-vector products through OrdinaryDiffEq's dense-output interpolation: SciML/SciMLSensitivity.jl#1648.
Runic wants spaces after commas and around <: in where clauses. The forward-mode test passed a NoTangent() for the closure, but inside @safetestset closures capture their free variables as real fields (unlike a plain top-level script), so it needs the closure's actual zero tangent instead.
Member
Author
|
@ChrisRackauckas this PR + chalk-lab/Mooncake.jl#1331 are only left to merge for fixing SciML/SciMLSensitivity.jl#1648. so HVP, plain reverse or forward mode all should work for that specifc issue. |
ChrisRackauckas
added a commit
that referenced
this pull request
Sep 22, 2026
Ships two new package extensions declaring the searches inactive for AD: EnzymeCore (#116) and Mooncake (#117). Minor because both add weak dependencies and new extension behaviour. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude <noreply@anthropic.com> Agent-Harness: Claude Code 2.1.273 Agent-Model: claude-opus-5[1m] Agent-Session: 6a01e420-3436-4518-83cc-6cd853283cfb (local transcript; this harness publishes no conversation URL)
ChrisRackauckas
added a commit
that referenced
this pull request
Sep 22, 2026
Ships two new package extensions declaring the searches inactive for AD: EnzymeCore (#116) and Mooncake (#117). Minor because both add weak dependencies and new extension behaviour. Agent-Harness: Claude Code 2.1.273 Agent-Model: claude-opus-5[1m] Agent-Session: 6a01e420-3436-4518-83cc-6cd853283cfb (local transcript; this harness publishes no conversation URL) Co-authored-by: Claude <noreply@anthropic.com>
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.
Mirrors the existing
EnzymeCoreextension:searchsorted_last,searchsorted_first,searchsortedlast!,searchsortedfirst!,searchsortedrange,findequal,findfirstequal, andfindfirstsortedequalall return a discrete index with no meaningful derivative, same asEnzymeRules.inactivealready marks them.For Mooncake this also works around a limitation:
KIND_SIMD_LINEAR_SCAN's kernel embeds a rawllvmcallSIMD intrinsic that Mooncake can't translate. Marking the dispatcher itself inactive stops Mooncake before it descends into that kernel at all, rather than needing a fix on Mooncake's side for this specific SIMD path.Found this while debugging native Hessian-vector products through
OrdinaryDiffEq's dense-output interpolation, which uses these functions to locate a query time in the sorted grid of saved solution times: SciML/SciMLSensitivity.jl#1648.Adds
ext/FindFirstFunctionsMooncakeExt.jl(zero-derivative rules for the plain signature and theCore.kwcallpath for keyword arguments likeorder/strategy/queries_sorted), theMooncakeweakdep/extension entries inProject.toml, andtest/Mooncake/mooncake_tests.jl(reverse mode, forward mode, native HVP, a dedicatedKIND_SIMD_LINEAR_SCANcase, and the batched/equality-search functions) wired intotest_groups.tomlas a new CI group.All 6 new test groups pass. Also Verified against the SciMLSensitivity HVP scenario using a Mooncake branch with chalk-lab/Mooncake.jl#1327 (solved via chalk-lab/Mooncake.jl#1329) and chalk-lab/Mooncake.jl#1331.