Make DifferentiateWith compatible with ForwardDiff, clarify docs (#487)

* Make DIfferentiateWith compatible with ForwardDiff, clarify docs

* More details

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/DifferentiationInterfaceForwardDiffExt.jl

@@ -7,6 +7,7 @@ import DifferentiationInterface as DI
 using DifferentiationInterface:
     Context,
     DerivativeExtras,
+    DifferentiateWith,
     GradientExtras,
     HessianExtras,
     HVPExtras,
@@ -59,5 +60,6 @@ include("utils.jl")
 include("onearg.jl")
 include("twoarg.jl")
 include("secondorder.jl")
+include("differentiate_with.jl")
 
 end # module

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/differentiate_with.jl

@@ -0,0 +1,15 @@
+function (dw::DifferentiateWith)(x::Dual{T,V,N}) where {T,V,N}
+    @compat (; f, backend) = dw
+    xval = myvalue(T, x)
+    tx = mypartials(T, Val(N), x)
+    y, ty = DI.value_and_pushforward(f, backend, xval, tx)
+    return make_dual(T, y, ty)
+end
+
+function (dw::DifferentiateWith)(x::AbstractArray{Dual{T,V,N}}) where {T,V,N}
+    @compat (; f, backend) = dw
+    xval = myvalue(T, x)
+    tx = mypartials(T, Val(N), x)
+    y, ty = DI.value_and_pushforward(f, backend, xval, tx)
+    return make_dual(T, y, ty)
+end

DifferentiationInterface/src/misc/differentiate_with.jl

@@ -1,18 +1,28 @@
 """
     DifferentiateWith
 
-Callable function wrapper that enforces differentiation with a specified (inner) backend.
+Function wrapper that enforces differentiation with a "substitute" AD backend, possible different from the "true" AD backend that is called.
 
-This works by defining new rules overriding the behavior of the outer backend that would normally be used.
+For instance, suppose a function `f` is not differentiable with Zygote because it involves mutation, but you know that it is differentiable with Enzyme.
+Then `f2 = DifferentiateWith(f, AutoEnzyme())` is a new function that behaves like `f`, except that `f2` is differentiable with Zygote (thanks to a chain rule which calls Enzyme under the hood).
+Moreover, any larger algorithm `alg` that calls `f2` instead of `f` will also be differentiable with Zygote (as long as `f` was the only Zygote blocker).
+
+!!! tip
+    This is mainly relevant for package developers who want to produce differentiable code at low cost, without writing the differentiation rules themselves.
+    If you sprinkle a few `DifferentiateWith` in places where some AD backends may struggle, end users can pick from a wider variety of packages to differentiate your algorithms.
 
 !!! warning
-    This is an experimental functionality, whose API cannot yet be considered stable.
-    It only supports out-of-place functions, and rules are only defined for [ChainRules.jl](https://github.com/JuliaDiff/ChainRules.jl)-compatible outer backends.
+    `DifferentiateWith` only supports out-of-place functions `y = f(x)` without additional context arguments.
+    It only makes these functions differentiable if the true backend is either [ForwardDiff](https://github.com/JuliaDiff/ForwardDiff.jl) or compatible with [ChainRules](https://github.com/JuliaDiff/ChainRules.jl).
+    For any other true backend, the differentiation behavior is not altered by `DifferentiateWith` (it becomes a transparent wrapper).
 
 # Fields
 
-- `f`: the function in question
-- `backend::AbstractADType`: the inner backend to use for differentiation
+- `f`: the function in question, with signature `f(x)`
+- `backend::AbstractADType`: the substitute backend to use for differentiation
+
+!!! note
+    For the substitute AD backend to be called under the hood, its package needs to be loaded in addition to the package of the true AD backend.
 
 # Constructor
 
@@ -21,40 +31,50 @@ This works by defining new rules overriding the behavior of the outer backend th
 # Example
 
 ```jldoctest
-using DifferentiationInterface
-import ForwardDiff, Zygote
+julia> using DifferentiationInterface
 
-function f(x)
-    a = Vector{eltype(x)}(undef, 1)
-    a[1] = sum(x)  # mutation that breaks Zygote
-    return a[1]
-end
+julia> import Enzyme, ForwardDiff, Zygote
+
+julia> function f(x::Vector{Float64})
+           a = Vector{Float64}(undef, 1)  # type constraint breaks ForwardDiff
+           a[1] = sum(abs2, x)  # mutation breaks Zygote
+           return a[1]
+       end;
 
-dw = DifferentiateWith(f, AutoForwardDiff());
+julia> f2 = DifferentiateWith(f, AutoEnzyme());
 
-gradient(dw, AutoZygote(), [2.0])  # calls ForwardDiff instead
+julia> f([3.0, 5.0]) == f2([3.0, 5.0])
+true
 
-# output
+julia> alg(x) = 7 * f2(x);
 
-1-element Vector{Float64}:
- 1.0
+julia> ForwardDiff.gradient(alg, [3.0, 5.0])
+2-element Vector{Float64}:
+ 42.0
+ 70.0
+
+julia> Zygote.gradient(alg, [3.0, 5.0])[1]
+2-element Vector{Float64}:
+ 42.0
+ 70.0
 ```
 """
 struct DifferentiateWith{F,B<:AbstractADType}
     f::F
     backend::B
 end
 
-"""
-    (dw::DifferentiateWith)(x)
-
-Call the underlying function `dw.f` of a [`DifferentiateWith`](@ref) wrapper.
-"""
 (dw::DifferentiateWith)(x) = dw.f(x)
 
 function Base.show(io::IO, dw::DifferentiateWith)
     @compat (; f, backend) = dw
     return print(
-        io, DifferentiateWith, "(", repr(f; context=io), ",", repr(backend; context=io), ")"
+        io,
+        DifferentiateWith,
+        "(",
+        repr(f; context=io),
+        ", ",
+        repr(backend; context=io),
+        ")",
     )
 end

DifferentiationInterface/test/Misc/DifferentiateWith/test.jl

@@ -1,31 +1,29 @@
 using Pkg
-Pkg.add(["ForwardDiff", "Zygote"])
+Pkg.add(["FiniteDiff", "ForwardDiff", "Zygote"])
 
 using DifferentiationInterface, DifferentiationInterfaceTest
 import DifferentiationInterfaceTest as DIT
+using FiniteDiff: FiniteDiff
 using ForwardDiff: ForwardDiff
 using Zygote: Zygote
 using Test
 
 LOGGING = get(ENV, "CI", "false") == "false"
 
-function zygote_breaking_scenarios()
-    outofplace_scens = filter(default_scenarios()) do scen
-        DIT.operator_place(scen) == :out
-    end
-    bad_outofplace_scens = map(outofplace_scens) do scen
-        function bad_f(x)
-            a = Vector{eltype(x)}(undef, 1)
-            a[1] = sum(x)
-            return scen.f(x)
+function differentiatewith_scenarios()
+    bad_scens =  # these closurified scenarios have mutation and type constraints
+        filter(default_scenarios(; include_normal=false, include_closurified=true)) do scen
+            DIT.function_place(scen) == :out
         end
-        wrapped_bad_f = DifferentiateWith(bad_f, AutoForwardDiff())
-        bad_scen = DIT.change_function(scen, wrapped_bad_f)
-        return bad_scen
+    good_scens = map(bad_scens) do scen
+        DIT.change_function(scen, DifferentiateWith(scen.f, AutoFiniteDiff()))
     end
-    return bad_outofplace_scens
+    return good_scens
 end
 
 test_differentiation(
-    AutoZygote(), zygote_breaking_scenarios(); second_order=false, logging=LOGGING
+    [AutoForwardDiff(), AutoZygote()],
+    differentiatewith_scenarios();
+    second_order=false,
+    logging=LOGGING,
 )