diff --git a/Project.toml b/Project.toml index 655c4a6..6b3c56c 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "FunctionWrappersWrappers" uuid = "77dc65aa-8811-40c2-897b-53d922fa7daf" authors = ["Chris Elrod and contributors"] -version = "1.9.0" +version = "1.9.1" [deps] FunctionWrappers = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e" diff --git a/ext/FunctionWrappersWrappersEnzymeExt.jl b/ext/FunctionWrappersWrappersEnzymeExt.jl index 42229f6..38097c5 100644 --- a/ext/FunctionWrappersWrappersEnzymeExt.jl +++ b/ext/FunctionWrappersWrappersEnzymeExt.jl @@ -1,10 +1,26 @@ module FunctionWrappersWrappersEnzymeExt using FunctionWrappersWrappers +using FunctionWrappersWrappers: SingleCacheStorage, DictCacheStorage, NoCacheStorage using Enzyme using EnzymeCore using EnzymeCore.EnzymeRules +# ============================================================================= +# Mark cache-storage types as inactive +# ============================================================================= +# `SingleCacheStorage` and `DictCacheStorage` are mutable / contain a `Dict`, +# and their cache-miss branches write to that storage. Without these +# declarations Enzyme conservatively treats any closure that *might* touch a +# `FunctionWrappersWrapper` (e.g. via `remake(prob; p = …)` capturing the +# problem in scope) as potentially writing to the wrapper's cache, and +# 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 + # ============================================================================= # Helper: build a Forward mode from FwdConfig flags # =============================================================================