QA: scan the Enzyme and Mooncake extensions

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

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"

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,
+            ),
+        ),
     ),
 )