Better docs

Author: gdalle

DifferentiationInterface/docs/make.jl

@@ -10,6 +10,7 @@ using Zygote: Zygote
 
 links = InterLinks(
     "ADTypes" => "https://sciml.github.io/ADTypes.jl/stable/",
+    "Reactant" => "https://enzymead.github.io/Reactant.jl/stable/",
     "SparseConnectivityTracer" => "https://adrianhill.de/SparseConnectivityTracer.jl/stable/",
     "SparseMatrixColorings" => "https://gdalle.github.io/SparseMatrixColorings.jl/stable/",
     "Symbolics" => "https://symbolics.juliasymbolics.org/stable/",

DifferentiationInterface/docs/src/api.md

@@ -144,3 +144,9 @@ DifferentiationInterface.AutoReverseFromPrimitive
 ```@docs
 DifferentiationInterface.Prep
 ```
+
+### Reactant
+
+```@docs
+DifferentiationInterface.to_reactant
+```

DifferentiationInterface/docs/src/explanation/backends.md

@@ -169,6 +169,10 @@ If a GTPSA [`Descriptor`](https://bmad-sim.github.io/GTPSA.jl/stable/man/b_descr
 
 Most operators fall back on `AutoForwardDiff`.
 
+### Reactant
+
+See the docstring for [`AutoReactant`](@ref).
+
 ### ReverseDiff
 
 With `AutoReverseDiff(compile=false)`, preparation preallocates a [config](https://juliadiff.org/ReverseDiff.jl/dev/api/#The-AbstractConfig-API).

DifferentiationInterface/ext/DifferentiationInterfaceReactantExt/onearg.jl

@@ -13,9 +13,9 @@ function DI.prepare_gradient_nokwarg(
     ) where {F, C}
     _sig = DI.signature(f, rebackend, x; strict)
     backend = rebackend.mode
-    xr = to_reac(x)
-    gr = to_reac(similar(x))
-    contextsr = map(to_reac, contexts)
+    xr = x isa ConcreteRArray ? nothing : ConcreteRArray(x)
+    gr = x isa ConcreteRArray ? nothing : ConcreteRArray(similar(x))
+    contextsr = map(_to_reactant, contexts)
     compiled_gradient = @compile DI.gradient(f, backend, xr, contextsr...)
     compiled_gradient! = @compile DI.gradient!(f, gr, backend, xr, contextsr...)
     compiled_value_and_gradient = @compile DI.value_and_gradient(f, backend, xr, contextsr...)
@@ -36,10 +36,9 @@ function DI.gradient(
     ) where {F, C}
     DI.check_prep(f, prep, rebackend, x)
     backend = rebackend.mode
-    (; xr, compiled_gradient) = prep
-    copyto!(xr, x)
-    contextsr = map(to_reac, contexts)
-    gr = compiled_gradient(f, backend, xr, contextsr...)
+    xr = isnothing(prep.xr) ? x : copyto!(prep.xr, x)
+    contextsr = map(_to_reactant, contexts)
+    gr = prep.compiled_gradient(f, backend, xr, contextsr...)
     return gr
 end
 
@@ -48,10 +47,9 @@ function DI.value_and_gradient(
     ) where {F, C}
     DI.check_prep(f, prep, rebackend, x)
     backend = rebackend.mode
-    (; xr, compiled_value_and_gradient) = prep
-    copyto!(xr, x)
-    contextsr = map(to_reac, contexts)
-    yr, gr = compiled_value_and_gradient(f, backend, xr, contextsr...)
+    xr = isnothing(prep.xr) ? x : copyto!(prep.xr, x)
+    contextsr = map(_to_reactant, contexts)
+    yr, gr = prep.compiled_value_and_gradient(f, backend, xr, contextsr...)
     return yr, gr
 end
 
@@ -60,21 +58,21 @@ function DI.gradient!(
     ) where {F, C}
     DI.check_prep(f, prep, rebackend, x)
     backend = rebackend.mode
-    (; xr, gr, compiled_gradient!) = prep
-    copyto!(xr, x)
-    contextsr = map(to_reac, contexts)
-    compiled_gradient!(f, gr, backend, xr, contextsr...)
-    return copyto!(grad, gr)
+    xr = isnothing(prep.xr) ? x : copyto!(prep.xr, x)
+    gr = isnothing(prep.gr) ? grad : prep.gr
+    contextsr = map(_to_reactant, contexts)
+    prep.compiled_gradient!(f, gr, backend, xr, contextsr...)
+    return isnothing(prep.gr) ? grad : copyto!(grad, gr)
 end
 
 function DI.value_and_gradient!(
         f::F, grad, prep::ReactantGradientPrep, rebackend::AutoReactant, x, contexts::Vararg{DI.Context, C}
     ) where {F, C}
     DI.check_prep(f, prep, rebackend, x)
     backend = rebackend.mode
-    (; xr, gr, compiled_value_and_gradient!) = prep
-    copyto!(xr, x)
-    contextsr = map(to_reac, contexts)
-    yr, gr = compiled_value_and_gradient!(f, gr, backend, xr, contextsr...)
-    return yr, copyto!(grad, gr)
+    xr = isnothing(prep.xr) ? x : copyto!(prep.xr, x)
+    gr = isnothing(prep.gr) ? grad : prep.gr
+    contextsr = map(_to_reactant, contexts)
+    yr, gr = prep.compiled_value_and_gradient!(f, gr, backend, xr, contextsr...)
+    return yr, isnothing(prep.gr) ? grad : copyto!(grad, gr)
 end

DifferentiationInterface/ext/DifferentiationInterfaceReactantExt/utils.jl

@@ -1,7 +1,3 @@
-to_reac(x::AbstractArray) = to_rarray(x)
-to_reac(x::ConcreteRArray) = x
-to_reac(x::Number) = ConcreteRNumber(x)
-to_reac(x::ConcreteRNumber) = x
-
-to_reac(c::DI.Constant) = DI.Constant(to_reac(DI.unwrap(c)))
-to_reac(c::DI.Cache) = DI.Cache(to_reac(DI.unwrap(c)))
+_to_reactant(x) = DI.to_reactant(x)
+_to_reactant(c::DI.Constant) = DI.Constant(_to_reactant(DI.unwrap(c)))
+_to_reactant(c::DI.Cache) = DI.Cache(_to_reactant(DI.unwrap(c)))

DifferentiationInterface/src/DifferentiationInterface.jl

@@ -69,6 +69,7 @@ include("misc/sparsity_detector.jl")
 include("misc/simple_finite_diff.jl")
 include("misc/zero_backends.jl")
 include("misc/overloading.jl")
+include("misc/reactant.jl")
 
 ## Exported
 
@@ -132,6 +133,7 @@ export AutoSparse
 @public inner, outer
 @public AutoForwardFromPrimitive, AutoReverseFromPrimitive
 @public Prep
+@public to_reactant
 
 include("init.jl")
 

DifferentiationInterface/src/misc/reactant.jl

@@ -0,0 +1,84 @@
+"""
+!!! tip "DI-specific information"
+    This part of the docstring is related to the use of `AutoReactant` inside [DifferentiationInterface.jl](https://github.com/JuliaDiff/DifferentiationInterface.jl), or DI for short.
+    Reactant's tutorial on [partial evaluation](https://enzymead.github.io/Reactant.jl/stable/tutorials/partial-evaluation) is useful reading to understand what follows.
+
+The `AutoReactant` backend inside DI imposes the following restrictions / assumptions:
+
+- The only supported operator (at the moment) is `DI.gradient` (along with its variants).
+- The input `x` must be an `AbstractArray` such that `Reactant.ConcreteRArray(x)` is well-defined.
+- By default, contexts such as `DI.Constant` and `DI.Cache` will be partially evaluated inside the compiled differentiation operator at preparation time. This means that the context value provided at preparation will be reused at every subsequent execution, while the context value provided at execution will be ignored. In particular, `DI.Cache` contexts will usually error and `DI.Constant` contexts will be frozen to one value.
+
+To disable partial evaluation and enforce tracing of contexts instead, first wrap them into types that _you own_.
+Then, overload [`DifferentiationInterface.to_reactant`](@ref) on these types to perform tracing in the way you see fit, for instance with `Reactant.to_rarray`.
+Every value you choose not to trace will still be partially evaluated at preparation time.
+
+# Example
+
+```jldoctest
+using DifferentiationInterface
+import DifferentiationInterface as DI
+import Reactant
+
+struct MyArgument{T1 <: Number, T2 <: AbstractArray}
+    u::T1
+    v::T2
+end
+
+f(x, a::MyArgument) = a.u * sum(a.v .* x .^ 2)
+
+DI.to_reactant(a::MyArgument) = Reactant.to_rarray(a; track_numbers = false)
+
+# preparation time
+x0 = zeros(2)
+a0 = MyArgument(1.0, [2.0, 3.0])
+
+# execution time
+x = [4.0, 5.0]
+a = MyArgument(6.0, [7.0, 8.0])
+
+backend = AutoReactant()
+prep = prepare_gradient(f, backend, x0, Constant(a0));
+
+g = gradient(f, prep, backend, x, Constant(a))
+g ≈ a0.u * 2 * (a.v .* x)  # a0.u is partially evaluated, a0.v is traced
+
+# output
+
+true
+```
+"""
+AutoReactant
+
+"""
+    to_reactant(a)
+
+Convert an argument `a` to an object `ar` containing the same values, where all the fields and subfields that can contain active (differentiated) data have been translated to [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl) types such as [`ConcreteRArray`](@extref Reactant.ConcreteRArray) or [`ConcreteRNumber`](@extref Reactant.ConcreteRNumber).
+
+!!! danger
+    DifferentiationInterface.jl implements this function as the identity, on purpose.
+    It should not be overloaded on base types, but only on types that you own, to modify the default behavior of `AutoReactant`.
+
+# Example
+
+```jldoctest
+import DifferentiationInterface as DI
+import Reactant
+
+struct MyArgument{T1 <: Number, T2 <: AbstractArray}
+    u::T1
+    v::T2
+end
+
+DI.to_reactant(a::MyArgument) = Reactant.to_rarray(a; track_numbers = false)
+
+a = MyArgument(1.0, [2.0, 3.0])
+ar = DI.to_reactant(a)
+ar isa MyArgument{Float64, <:Reactant.ConcreteRArray}
+
+# output
+
+true
+```
+"""
+to_reactant(x) = x

DifferentiationInterface/test/Back/Reactant/Project.toml

@@ -0,0 +1,10 @@
+[deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
+DifferentiationInterfaceTest = "a82114a7-5aa3-49a8-9643-716bb13727a3"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
+SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

DifferentiationInterface/test/Back/Reactant/test.jl

@@ -1,20 +1,24 @@
-using Pkg
-Pkg.add(url = "https://github.com/EnzymeAD/Enzyme.jl")
-Pkg.add("Reactant")
+include("../../testutils.jl")
 
 using DifferentiationInterface
+import DifferentiationInterface as DI
 using DifferentiationInterfaceTest
-using Reactant
+import DifferentiationInterfaceTest as DIT
+import Enzyme, Reactant
 using Test
 
 backend = AutoReactant()
 
 @test check_available(backend)
 @test check_inplace(backend)
 
+scen1 = DIT.Scenario(
+
+)
+
 test_differentiation(
     backend, DifferentiationInterfaceTest.default_scenarios(;
-        include_constantified = true, include_cachified = false
+        include_constantified = false, include_cachified = false
     );
     excluded = vcat(SECOND_ORDER, :jacobian, :derivative, :pushforward, :pullback),
     logging = false