Start implementing Cache contexts (#587)

* Start implementing caches

* First Cache implementation and tests

* Function modifiers

* Scenario conversion

Author: gdalle

DifferentiationInterface/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.6.14"
+version = "0.6.15"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

DifferentiationInterface/docs/src/api.md

@@ -13,6 +13,7 @@ DifferentiationInterface
 ```@docs
 Context
 Constant
+Cache
 ```
 
 ## First order

DifferentiationInterface/docs/src/explanation/advanced.md

@@ -10,22 +10,29 @@ However, the release v0.6 introduced the possibility of additional "context" arg
 Contexts can be useful if you have a function `y = f(x, a, b, c, ...)` or `f!(y, x, a, b, c, ...)` and you want derivatives of `y` with respect to `x` only.
 Another option would be creating a closure, but that is sometimes undesirable.
 
-!!! warning
-    This feature is still experimental, and will likely not be supported by all backends.
-    At the moment, it only works with certain backends, among which ForwardDiff, Zygote and Enzyme.
-
 ### Types of contexts
 
 Every context argument must be wrapped in a subtype of [`Context`](@ref) and come after the differentiated input `x`.
-Right now, there is only one kind of context, namely [`Constant`](@ref), but we might add more.
-Semantically, calling
+Right now, there are two kinds of context: [`Constant`](@ref) and [`Cache`](@ref).
+
+!!! warning
+    This feature is still experimental and will not be supported by all backends.
+    At the moment:
+    - `Constant` is supported by all backends except symbolic ones
+    - `Cache` is only supported by finite difference backends
+
+Semantically, both of these calls compute the partial gradient of `f(x, c)` with respect to `x`, but they consider `c` differently:
 
 ```julia
 gradient(f, backend, x, Constant(c))
+gradient(f, backend, x, Cache(c))
 ```
 
-computes the partial gradient of `f(x, c)` with respect to `x`, while keeping `c` constant.
-Importantly, one can prepare an operator with an arbitrary value `c'` of the constant (subject to the usual restrictions on preparation).
+In the first call, `c` is kept unchanged throughout the function evaluation.
+In the second call, `c` can be mutated with values computed during the function.
+
+Importantly, one can prepare an operator with an arbitrary value `c'` of the `Constant` (subject to the usual restrictions on preparation).
+The values in a provided `Cache` never matter anyway.
 
 ## Sparsity
 

DifferentiationInterface/src/DifferentiationInterface.jl

@@ -65,7 +65,7 @@ include("misc/zero_backends.jl")
 
 ## Exported
 
-export Context, Constant
+export Context, Constant, Cache
 export SecondOrder
 
 export value_and_pushforward!, value_and_pushforward

DifferentiationInterface/src/utils/context.jl

@@ -15,6 +15,7 @@ Abstract supertype for additional context arguments, which can be passed to diff
 # See also
 
 - [`Constant`](@ref)
+- [`Cache`](@ref)
 """
 abstract type Context end
 
@@ -58,6 +59,25 @@ end
 
 Base.convert(::Type{Constant{T}}, x) where {T} = Constant(convert(T, x))
 
+"""
+    Cache
+
+Concrete type of [`Context`](@ref) argument which can be mutated with active values during differentiation.
+
+The initial values present inside the cache do not matter.
+"""
+struct Cache{T} <: Context
+    data::T
+end
+
+unwrap(c::Cache) = c.data
+
+function Base.convert(::Type{Cache{T}}, x::Cache) where {T}
+    return Cache(convert(T, x.data))
+end
+
+Base.convert(::Type{Cache{T}}, x) where {T} = Cache(convert(T, x))
+
 struct Rewrap{C,T}
     function Rewrap(contexts::Vararg{Context,C}) where {C}
         T = typeof(contexts)

DifferentiationInterface/test/Back/FiniteDiff/test.jl

@@ -14,7 +14,7 @@ end
 
 test_differentiation(
     AutoFiniteDiff(),
-    default_scenarios(; include_constantified=true);
+    default_scenarios(; include_constantified=true, include_cachified=true);
     excluded=[:second_derivative, :hvp],
     logging=LOGGING,
 );

DifferentiationInterface/test/Back/FiniteDifferences/test.jl

@@ -14,7 +14,7 @@ end
 
 test_differentiation(
     AutoFiniteDifferences(; fdm=FiniteDifferences.central_fdm(3, 1)),
-    default_scenarios(; include_constantified=true);
+    default_scenarios(; include_constantified=true, include_cachified=true);
     excluded=SECOND_ORDER,
     logging=LOGGING,
 );

DifferentiationInterfaceTest/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterfaceTest"
 uuid = "a82114a7-5aa3-49a8-9643-716bb13727a3"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.8.3"
+version = "0.8.4"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

DifferentiationInterfaceTest/ext/DifferentiationInterfaceTestJLArraysExt/DifferentiationInterfaceTestJLArraysExt.jl

@@ -21,13 +21,13 @@ function (f::DIT.NumToArr{JLMatrix{T}})(x::Number) where {T}
     return sin.(x .* a)
 end
 
-myjl(f::DIT.MultiplyByConstant) = f
-myjl(f::DIT.WritableClosure) = f
+myjl(f::DIT.FunctionModifier) = f
 
 myjl(x::Number) = x
 myjl(x::AbstractArray) = jl(x)
 myjl(x::Tuple) = map(myjl, x)
 myjl(x::DI.Constant) = DI.Constant(myjl(DI.unwrap(x)))
+myjl(x::DI.Cache) = DI.Cache(myjl(DI.unwrap(x)))
 myjl(::Nothing) = nothing
 
 function myjl(scen::Scenario{op,pl_op,pl_fun}) where {op,pl_op,pl_fun}

DifferentiationInterfaceTest/ext/DifferentiationInterfaceTestStaticArraysExt/DifferentiationInterfaceTestStaticArraysExt.jl

@@ -7,45 +7,45 @@ using Random: AbstractRNG, default_rng
 using SparseArrays: SparseArrays, SparseMatrixCSC, nnz, spdiagm
 using StaticArrays: MArray, MMatrix, MVector, SArray, SMatrix, SVector
 
-mySArray(f::Function) = f
-mySArray(::DIT.NumToArr{A}) where {T,A<:AbstractVector{T}} = DIT.NumToArr(SVector{6,T})
-mySArray(::DIT.NumToArr{A}) where {T,A<:AbstractMatrix{T}} = DIT.NumToArr(SMatrix{2,3,T,6})
-mySArray(f::DIT.MultiplyByConstant) = f
-mySArray(f::DIT.WritableClosure) = f
+mystatic(f::Function) = f
+mystatic(::DIT.NumToArr{A}) where {T,A<:AbstractVector{T}} = DIT.NumToArr(SVector{6,T})
+mystatic(::DIT.NumToArr{A}) where {T,A<:AbstractMatrix{T}} = DIT.NumToArr(SMatrix{2,3,T,6})
+mystatic(f::DIT.FunctionModifier) = f
 
-mySArray(x::Number) = x
-myMArray(x::Number) = x
+mystatic(x::Number) = x
+mymutablestatic(x::Number) = x
 
-mySArray(x::AbstractVector{T}) where {T} = convert(SVector{length(x),T}, x)
-myMArray(x::AbstractVector{T}) where {T} = convert(MVector{length(x),T}, x)
+mystatic(x::AbstractVector{T}) where {T} = convert(SVector{length(x),T}, x)
+mymutablestatic(x::AbstractVector{T}) where {T} = convert(MVector{length(x),T}, x)
 
-function mySArray(x::AbstractMatrix{T}) where {T}
+function mystatic(x::AbstractMatrix{T}) where {T}
     return convert(SMatrix{size(x, 1),size(x, 2),T,length(x)}, x)
 end
-function myMArray(x::AbstractMatrix{T}) where {T}
+function mymutablestatic(x::AbstractMatrix{T}) where {T}
     return convert(MMatrix{size(x, 1),size(x, 2),T,length(x)}, x)
 end
 
-mySArray(x::Tuple) = map(mySArray, x)
-mySArray(x::DI.Constant) = DI.Constant(mySArray(DI.unwrap(x)))
-mySArray(::Nothing) = nothing
+mystatic(x::Tuple) = map(mystatic, x)
+mystatic(x::DI.Constant) = DI.Constant(mystatic(DI.unwrap(x)))
+mystatic(x::DI.Cache) = DI.Cache(mymutablestatic(DI.unwrap(x)))
+mystatic(::Nothing) = nothing
 
-function mySArray(scen::Scenario{op,pl_op,pl_fun}) where {op,pl_op,pl_fun}
+function mystatic(scen::Scenario{op,pl_op,pl_fun}) where {op,pl_op,pl_fun}
     (; f, x, y, tang, contexts, res1, res2) = scen
     return Scenario{op,pl_op,pl_fun}(
-        mySArray(f);
-        x=mySArray(x),
-        y=pl_fun == :in ? myMArray(y) : mySArray(y),
-        tang=mySArray(tang),
-        contexts=mySArray(contexts),
-        res1=mySArray(res1),
-        res2=mySArray(res2),
+        mystatic(f);
+        x=mystatic(x),
+        y=pl_fun == :in ? mymutablestatic(y) : mystatic(y),
+        tang=mystatic(tang),
+        contexts=mystatic(contexts),
+        res1=mystatic(res1),
+        res2=mystatic(res2),
     )
 end
 
 function DIT.static_scenarios(args...; kwargs...)
     scens = DIT.default_scenarios(args...; kwargs...)
-    return mySArray.(scens)
+    return mystatic.(scens)
 end
 
 end