From 96875e4100beee0566326cc771376750fd0d72cf Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Sat, 1 Aug 2026 04:51:33 -0400 Subject: [PATCH] Scan the Enzyme and Mooncake extensions in QA ExplicitImports only analyzes an extension module when `Base.get_extension` returns it, which requires the extension's triggers to be loaded. The QA environment loaded no weakdeps, so neither `FunctionWrappersWrappersEnzymeExt` nor `FunctionWrappersWrappersMooncakeExt` was checked at all. Add Enzyme, EnzymeCore and Mooncake to `test/qa/Project.toml` (compat mirroring the root Project.toml) and `using` them in `qa.jl`. Both extensions are now covered. The new findings in the Enzyme extension are fixed at the source rather than ignored: the blanket `using Enzyme` / `using EnzymeCore` / `using EnzymeCore.EnzymeRules` become explicit imports, and the `inactive_type` definitions go through `EnzymeRules` directly instead of the `EnzymeCore.EnzymeRules` path. This also drops the two "both EnzymeRules and Enzyme export ..." ambiguity warnings emitted when the extension loaded. The remaining findings are AD rule-authoring interfaces with no public spelling (EnzymeRules' extension points, Mooncake's rule API) plus `Core.Typeof`, and are ignored with a comment naming the owner. Co-Authored-By: Chris Rackauckas --- ext/FunctionWrappersWrappersEnzymeExt.jl | 12 +++---- test/qa/Project.toml | 6 ++++ test/qa/qa.jl | 41 +++++++++++++++++++++++- 3 files changed, 52 insertions(+), 7 deletions(-) diff --git a/ext/FunctionWrappersWrappersEnzymeExt.jl b/ext/FunctionWrappersWrappersEnzymeExt.jl index b401cf8..ef773fb 100644 --- a/ext/FunctionWrappersWrappersEnzymeExt.jl +++ b/ext/FunctionWrappersWrappersEnzymeExt.jl @@ -2,9 +2,9 @@ module FunctionWrappersWrappersEnzymeExt using FunctionWrappersWrappers using FunctionWrappersWrappers: SingleCacheStorage, DictCacheStorage, NoCacheStorage -using Enzyme -using EnzymeCore -using EnzymeCore.EnzymeRules +using Enzyme: Enzyme +using EnzymeCore: EnzymeCore, BatchDuplicated, Const, Duplicated, Forward, ForwardWithPrimal +using EnzymeCore.EnzymeRules: EnzymeRules # ============================================================================= # Mark cache-storage types as inactive @@ -17,9 +17,9 @@ using EnzymeCore.EnzymeRules # refuses to prove the captured argument read-only. The cache values are # `FunctionWrapper`s used purely for dispatch / dynamic call speedup; they # never hold derivative data. -EnzymeCore.EnzymeRules.inactive_type(::Type{<:SingleCacheStorage}) = true -EnzymeCore.EnzymeRules.inactive_type(::Type{<:DictCacheStorage}) = true -EnzymeCore.EnzymeRules.inactive_type(::Type{NoCacheStorage}) = true +EnzymeRules.inactive_type(::Type{<:SingleCacheStorage}) = true +EnzymeRules.inactive_type(::Type{<:DictCacheStorage}) = true +EnzymeRules.inactive_type(::Type{NoCacheStorage}) = true # ============================================================================= # Helper: build a Forward mode from FwdConfig flags diff --git a/test/qa/Project.toml b/test/qa/Project.toml index 597402f..368ca8c 100644 --- a/test/qa/Project.toml +++ b/test/qa/Project.toml @@ -1,14 +1,20 @@ [deps] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" FunctionWrappersWrappers = "77dc65aa-8811-40c2-897b-53d922fa7daf" JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" SciMLTesting = "09d9d899-5365-40a9-917a-5f67fddea283" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [compat] Aqua = "0.8" +Enzyme = "0.13.185" +EnzymeCore = "0.8" JET = "0.9, 0.10, 0.11" +Mooncake = "0.5" SafeTestsets = "0.0.1, 0.1, 1" SciMLTesting = "2.4" Test = "1" diff --git a/test/qa/qa.jl b/test/qa/qa.jl index 334f02e..9f853a1 100644 --- a/test/qa/qa.jl +++ b/test/qa/qa.jl @@ -1,5 +1,11 @@ using SciMLTesting, FunctionWrappersWrappers +# ExplicitImports only sees an extension module once its triggers are loaded +# (`Base.get_extension` returns `nothing` otherwise), so loading the weakdeps here +# is what puts `FunctionWrappersWrappersEnzymeExt` and +# `FunctionWrappersWrappersMooncakeExt` under QA at all. +using Enzyme, EnzymeCore, Mooncake + run_qa( FunctionWrappersWrappers; ei_kwargs = (; @@ -10,6 +16,39 @@ run_qa( # construction path, so it cannot be hidden behind our own API. Drop this # entry once JuliaLang/FunctionWrappers.jl#35 is released and the compat # floor is raised. - all_qualified_accesses_are_public = (; ignore = (:FunctionWrapper,)), + # + # `Mooncake` declares only its top-level differentiation entry points public + # (`value_and_gradient!!`, `prepare_gradient_cache`, `Config`, `Dual`, ...). + # Its whole rule-authoring interface — the thing an extension that teaches + # Mooncake about a new callable has to use — is documented but not `public`: + # `build_rrule`, `rrule!!`, `tangent_type` and `primal` below, plus the + # explicitly-imported names in the matching list. There is no public spelling + # for writing a Mooncake rule, so these stay ignored until Mooncake marks its + # rule API `public`. + # + # `Core.Typeof` is the standard idiom for building a call-signature tuple type + # (`typeof` is wrong for arguments that are themselves types); `Core` does not + # declare it public and there is no `Base` equivalent. + # + # `EnzymeCore.EnzymeRules` exports the config accessors a rule reads + # (`needs_primal`, `overwritten`, `runtime_activity`, ...) but not the generic + # functions a rule author adds methods to. `forward`, `augmented_primal`, + # `reverse`, `inactive_type` and `strong_zero` are exactly the documented + # extension points of the EnzymeRules interface, and there is no public + # spelling for them. + all_qualified_accesses_are_public = (; + ignore = ( + :FunctionWrapper, + :augmented_primal, :forward, :inactive_type, :reverse, :strong_zero, + :build_rrule, :primal, Symbol("rrule!!"), :tangent_type, + :Typeof, + ), + ), + all_explicit_imports_are_public = (; + ignore = ( + Symbol("@is_primitive"), :CoDual, :MinimalCtx, :NoRData, :NoTangent, + :fdata, :zero_tangent, + ), + ), ), )