Recursive caches

Author: gdalle

.github/workflows/Test.yml

@@ -25,7 +25,7 @@ jobs:
       actions: write
       contents: read
     strategy:
-      fail-fast: true  # TODO: toggle
+      fail-fast: false  # TODO: toggle
       matrix:
         version:
           - "1.10"

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/utils.jl

@@ -53,9 +53,7 @@ force_annotation(f::F) where {F} = Const(f)
     return Const(DI.unwrap(c))
 end
 
-@inline function _translate(
-    backend::AutoEnzyme, ::Mode, ::Val{B}, c::DI.GeneralizedCache
-) where {B}
+@inline function _translate(backend::AutoEnzyme, ::Mode, ::Val{B}, c::DI.Cache) where {B}
     if B == 1
         return Duplicated(DI.unwrap(c), make_zero(DI.unwrap(c)))
     else

DifferentiationInterface/ext/DifferentiationInterfaceFastDifferentiationExt/DifferentiationInterfaceFastDifferentiationExt.jl

@@ -23,6 +23,11 @@ myvec(x::AbstractArray) = vec(x)
 
 variablize(::Number, name::Symbol) = only(make_variables(name))
 variablize(x::AbstractArray, name::Symbol) = make_variables(name, size(x)...)
+function variablize(x::Union{Tuple,NamedTuple}, name::Symbol)
+    return map(x) do xi
+        variablize(xi, gensym())  # TODO: fix symbol?
+    end
+end
 
 function variablize(contexts::NTuple{C,DI.Context}) where {C}
     map(enumerate(contexts)) do (k, c)

DifferentiationInterface/ext/DifferentiationInterfaceSparseConnectivityTracerExt/DifferentiationInterfaceSparseConnectivityTracerExt.jl

@@ -5,26 +5,23 @@ import DifferentiationInterface as DI
 using SparseConnectivityTracer:
     TracerSparsityDetector, TracerLocalSparsityDetector, jacobian_buffer, hessian_buffer
 
-@inline _jacobian_translate(detector, c::DI.Constant) = DI.unwrap(c)
-@inline function _jacobian_translate(detector, c::DI.Cache{<:AbstractArray})
-    return jacobian_buffer(DI.unwrap(c), detector)
+@inline _translate(::Type, c::DI.Constant) = DI.unwrap(c)
+@inline function _translate(::Type{T}, c::DI.Cache) where {T}
+    return DI.recursive_similar(DI.unwrap(c), T)
 end
 
-function jacobian_translate(detector, contexts::Vararg{DI.Context,C}) where {C}
+function jacobian_translate(detector, x, contexts::Vararg{DI.Context,C}) where {C}
+    T = eltype(jacobian_buffer(x, detector))
     new_contexts = map(contexts) do c
-        _jacobian_translate(detector, c)
+        _translate(T, c)
     end
     return new_contexts
 end
 
-@inline _hessian_translate(detector, c::DI.Constant) = DI.unwrap(c)
-@inline function _hessian_translate(detector, c::DI.Cache{<:AbstractArray})
-    return hessian_buffer(DI.unwrap(c), detector)
-end
-
-function hessian_translate(detector, contexts::Vararg{DI.Context,C}) where {C}
+function hessian_translate(detector, x, contexts::Vararg{DI.Context,C}) where {C}
+    T = eltype(hessian_buffer(x, detector))
     new_contexts = map(contexts) do c
-        _hessian_translate(detector, c)
+        _translate(T, c)
     end
     return new_contexts
 end
@@ -35,7 +32,7 @@ function DI.jacobian_sparsity_with_contexts(
     x,
     contexts::Vararg{DI.Context,C},
 ) where {F,C}
-    contexts_tracer = jacobian_translate(detector, contexts...)
+    contexts_tracer = jacobian_translate(detector, x, contexts...)
     fc = DI.FixTail(f, contexts_tracer...)
     return jacobian_sparsity(fc, x, detector)
 end
@@ -47,7 +44,7 @@ function DI.jacobian_sparsity_with_contexts(
     x,
     contexts::Vararg{DI.Context,C},
 ) where {F,C}
-    contexts_tracer = jacobian_translate(detector, contexts...)
+    contexts_tracer = jacobian_translate(detector, x, contexts...)
     fc! = DI.FixTail(f!, contexts_tracer...)
     return jacobian_sparsity(fc!, y, x, detector)
 end

DifferentiationInterface/ext/DifferentiationInterfaceZygoteExt/DifferentiationInterfaceZygoteExt.jl

@@ -17,7 +17,10 @@ DI.check_available(::AutoZygote) = true
 DI.inplace_support(::AutoZygote) = DI.InPlaceNotSupported()
 
 translate(c::DI.Context) = DI.unwrap(c)
-translate(c::DI.Cache) = Buffer(DI.unwrap(c))
+translate(c::DI.Cache{<:AbstractArray}) = Buffer(DI.unwrap(c))
+function translate(c::DI.Cache{<:Union{NTuple,NamedTuple}})
+    return map(translate, map(DI.Cache, DI.unwrap(c)))
+end
 
 ## Pullback
 

DifferentiationInterface/src/utils/context.jl

@@ -23,7 +23,6 @@ Abstract supertype for additional context arguments, which can be passed to diff
 abstract type Context end
 
 abstract type GeneralizedConstant <: Context end
-abstract type GeneralizedCache <: Context end
 
 unwrap(c::Context) = c.data
 Base.:(==)(c1::Context, c2::Context) = unwrap(c1) == unwrap(c2)
@@ -78,7 +77,7 @@ The initial values present inside the cache do not matter.
 For some backends, preparation allocates the required memory for `Cache` contexts with the right element type, similar to [PreallocationTools.jl](https://github.com/SciML/PreallocationTools.jl).
 
 !!! warning
-    Most backends require any `Cache` context to be an `AbstractArray`.
+    Some backends require any `Cache` context to be an `AbstractArray or a (named) tuple of `AbstractArray`s.
 
 # Example
 
@@ -97,7 +96,7 @@ julia> gradient(f, prep, AutoForwardDiff(), [3.0, 4.0], Cache(zeros(2)))
  1.0
 ````
 """
-struct Cache{T} <: GeneralizedCache
+struct Cache{T} <: Context
     data::T
 end
 
@@ -114,12 +113,10 @@ struct BackendContext{T} <: GeneralizedConstant
     data::T
 end
 
-struct PrepContext{T} <: GeneralizedCache
+struct PrepContext{T} <: Context
     data::T
 end
 
-struct UnknownContext <: Context end
-
 ## Context manipulation
 
 struct Rewrap{C,T}
@@ -146,4 +143,4 @@ function with_contexts(f::F, contexts::Vararg{Context,N}) where {F,N}
 end
 
 adapt_eltype(c::Constant, ::Type) = c
-adapt_eltype(c::Cache, ::Type{T}) where {T} = Cache(similar(unwrap(c), T))
+adapt_eltype(c::Cache, ::Type{T}) where {T} = Cache(recursive_similar(unwrap(c), T))

DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl

@@ -7,6 +7,12 @@ using DifferentiationInterface:
 using SparseMatrixColorings
 using Test
 
+test_differentiation(
+    AutoSimpleFiniteDiff(),
+    default_scenarios(; include_normal=false, include_cachified=true);
+    logging=true,
+)
+
 LOGGING = get(ENV, "CI", "false") == "false"
 
 backends = [ #

DifferentiationInterfaceTest/src/scenarios/modify.jl

@@ -147,7 +147,7 @@ end
 """
     constantify(scen::Scenario)
 
-Return a new `Scenario` identical to `scen` except for the function `f`, which is made to accept an additional constant argument `a` by which the output is multiplied.
+Return a new `Scenario` identical to `scen` except for the function `f`, which is made to accept an additional constant argument by which the output is multiplied.
 The output and result fields are updated accordingly.
 """
 function constantify(scen::Scenario{op,pl_op,pl_fun}) where {op,pl_op,pl_fun}
@@ -178,7 +178,8 @@ end
 
 Base.show(io::IO, f::StoreInCache) = print(io, "StoreInCache($(f.f))")
 
-function (sc::StoreInCache{:out})(x, y_cache)
+function (sc::StoreInCache{:out})(x, y_cache_tup)
+    y_cache = y_cache_tup.cache
     y = sc.f(x)
     if y isa Number
         y_cache[1] = y
@@ -189,7 +190,8 @@ function (sc::StoreInCache{:out})(x, y_cache)
     end
 end
 
-function (sc::StoreInCache{:in})(y, x, y_cache)
+function (sc::StoreInCache{:in})(y, x, y_cache_tup)
+    y_cache = y_cache_tup.cache
     sc.f(y_cache, x)
     copyto!(y, y_cache)
     return nothing
@@ -198,16 +200,18 @@ end
 """
     cachify(scen::Scenario)
 
-Return a new `Scenario` identical to `scen` except for the function `f`, which is made to accept an additional cache argument `a` to store the result before it is returned.
+Return a new `Scenario` identical to `scen` except for the function `f`, which is made to accept an additional cache argument to store the result before it is returned.
+
+If `tup=true` the cache is a tuple of arrays, otherwise just an array.
 """
 function cachify(scen::Scenario{op,pl_op,pl_fun}) where {op,pl_op,pl_fun}
     (; f,) = scen
     @assert isempty(scen.contexts)
     cache_f = StoreInCache{pl_fun}(f)
     y_cache = if scen.y isa Number
-        [myzero(scen.y)]
+        (; cache=[myzero(scen.y)])
     else
-        mysimilar(scen.y)
+        (; cache=mysimilar(scen.y))
     end
     return Scenario{op,pl_op,pl_fun}(
         cache_f;