From 323c05a60f2194f2b2a3b5621a399fde147175ea Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Thu, 21 May 2026 11:20:27 -0400 Subject: [PATCH] Declare cache-storage types as Enzyme inactive MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `SingleCacheStorage` is a mutable struct and `DictCacheStorage` carries a mutable `Dict`. Their fallback paths write a `FunctionWrapper` into the storage on a cache miss. Without an `inactive_type` declaration Enzyme conservatively treats any closure that captures a `FunctionWrappersWrapper` as potentially writing to the storage, so `Enzyme.gradient`/`autodiff` through such a closure can fail with `EnzymeMutabilityException` on the captured argument. The cache values are `FunctionWrapper`s — used purely for dispatch and dynamic-call speedup. They never carry derivative data. Marking the three storage types as `inactive_type` lets Enzyme prove the captured wrapper readonly without affecting derivative correctness. Defensive hardening prompted by the Enzyme-through-MTK-remake stack documented in SciML/SciMLSensitivity.jl#1323, SciML/SciMLSensitivity.jl#1359, and SciML/ModelingToolkit.jl#4550 — this PR alone is not sufficient to unblock that path (the underlying NonlinearSolve + Enzyme interaction still hits a separate `EnzymeMutabilityException` against the wrapped function), but it removes one conservative-IPA layer for any user of FWW under Enzyme. Co-Authored-By: Chris Rackauckas --- Project.toml | 2 +- ext/FunctionWrappersWrappersEnzymeExt.jl | 16 ++++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) 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 # =============================================================================