Implement DifferentiateWith to translate between backends (#218)

* Implement DifferentiateWith to translate between backends

* Fix parsing

Author: gdalle

DifferentiationInterface/docs/src/api.md

@@ -91,6 +91,12 @@ check_twoarg
 check_hessian
 ```
 
+## Translation
+
+```@docs
+DifferentiateWith
+```
+
 ## Internals
 
 This is not part of the public API.

DifferentiationInterface/docs/src/backends.md

@@ -5,7 +5,7 @@ CollapsedDocStrings = true
 
 ```@setup backends
 using DifferentiationInterface
-using DifferentiationInterface: backend_string
+using DifferentiationInterface: backend_str
 import Markdown
 import Diffractor, Enzyme, FastDifferentiation, FiniteDiff, FiniteDifferences, ForwardDiff, PolyesterForwardDiff, ReverseDiff, Tapir, Tracker, Zygote
 
@@ -37,7 +37,7 @@ println(io, "|:--------|:------------:|:----------------------:|:---------------
 
 for example in backend_examples
     b = eval(Meta.parse(example)) # backend
-    join(io, [backend_string(b), unicode_check_available(b), unicode_check_twoarg(b), unicode_check_hessian(b), "`$example`"], '|')
+    join(io, [backend_str(b), unicode_check_available(b), unicode_check_twoarg(b), unicode_check_hessian(b), "`$example`"], '|')
     println(io, '|' )
 end
 backend_table = Markdown.parse(String(take!(io)))

DifferentiationInterface/docs/src/overloads.md

@@ -24,7 +24,7 @@ Each cell can have three values:
 ```@setup overloads
 using ADTypes: AbstractADType
 using DifferentiationInterface
-using DifferentiationInterface: backend_string, mutation_support, MutationSupported
+using DifferentiationInterface: backend_str, mutation_support, MutationSupported
 using Markdown: Markdown
 using Diffractor: Diffractor
 using Enzyme: Enzyme

DifferentiationInterface/docs/src/overview.md

@@ -123,6 +123,12 @@ We make this available for all backends with the following operators:
 | :--------------------------------- | :---------------------------------- |
 | [`value_and_pullback_split`](@ref) | [`value_and_pullback!_split`](@ref) |
 
+## Translation
+
+The wrapper [`DifferentiateWith`](@ref) allows you to take a function and specify that it should be differentiated with the backend of your choice.
+In other words, when you try to differentiate `dw = DifferentiateWith(f, backend1)` with `backend2`, then `backend1` steps in and `backend2` does nothing.
+At the moment it only works when `backend2` supports [ChainRules.jl](https://github.com/JuliaDiff/ChainRules.jl).
+
 ## Going further
 
 ### Non-standard types

DifferentiationInterface/ext/DifferentiationInterfaceChainRulesCoreExt/DifferentiationInterfaceChainRulesCoreExt.jl

@@ -2,9 +2,15 @@ module DifferentiationInterfaceChainRulesCoreExt
 
 using ADTypes: ADTypes, AutoChainRules
 using ChainRulesCore:
-    HasForwardsMode, HasReverseMode, NoTangent, RuleConfig, frule_via_ad, rrule_via_ad
+    ChainRulesCore,
+    HasForwardsMode,
+    HasReverseMode,
+    NoTangent,
+    RuleConfig,
+    frule_via_ad,
+    rrule_via_ad
 import DifferentiationInterface as DI
-using DifferentiationInterface: NoPullbackExtras, NoPushforwardExtras
+using DifferentiationInterface: DifferentiateWith, NoPullbackExtras, NoPushforwardExtras
 
 ruleconfig(backend::AutoChainRules) = backend.ruleconfig
 
@@ -14,32 +20,7 @@ const AutoReverseChainRules = AutoChainRules{<:RuleConfig{>:HasReverseMode}}
 DI.check_available(::AutoChainRules) = true
 DI.mutation_support(::AutoChainRules) = DI.MutationNotSupported()
 
-## Pullback
-
-DI.prepare_pullback(f, ::AutoReverseChainRules, x, dy) = NoPullbackExtras()
-
-function DI.value_and_pullback_split(
-    f, backend::AutoReverseChainRules, x, ::NoPullbackExtras
-)
-    rc = ruleconfig(backend)
-    y, pullback = rrule_via_ad(rc, f, x)
-    pullbackfunc(dy) = last(pullback(dy))
-    return y, pullbackfunc
-end
-
-function DI.value_and_pullback!_split(
-    f, backend::AutoReverseChainRules, x, extras::NoPullbackExtras
-)
-    y, pullbackfunc = DI.value_and_pullback_split(f, backend, x, extras)
-    pullbackfunc!(dx, dy) = copyto!(dx, pullbackfunc(dy))
-    return y, pullbackfunc!
-end
-
-function DI.value_and_pullback(
-    f, backend::AutoReverseChainRules, x, dy, extras::NoPullbackExtras
-)
-    y, pullbackfunc = DI.value_and_pullback_split(f, backend, x, extras)
-    return y, pullbackfunc(dy)
-end
+include("reverse_onearg.jl")
+include("differentiate_with.jl")
 
 end

DifferentiationInterface/ext/DifferentiationInterfaceChainRulesCoreExt/differentiate_with.jl

@@ -0,0 +1,12 @@
+function ChainRulesCore.frule((_, dx), dw::DifferentiateWith, x)
+    (; f, backend) = dw
+    y, dy = DI.value_and_pushforward(f, backend, x, dx)
+    return y, dy
+end
+
+function ChainRulesCore.rrule(dw::DifferentiateWith, x)
+    (; f, backend) = dw
+    y, pullbackfunc = DI.value_and_pullback_split(f, backend, x)
+    pullbackfunc_adjusted(dy) = (NoTangent(), pullbackfunc(dy))
+    return y, pullbackfunc_adjusted
+end

DifferentiationInterface/ext/DifferentiationInterfaceChainRulesCoreExt/reverse_onearg.jl

@@ -0,0 +1,27 @@
+## Pullback
+
+DI.prepare_pullback(f, ::AutoReverseChainRules, x, dy) = NoPullbackExtras()
+
+function DI.value_and_pullback_split(
+    f, backend::AutoReverseChainRules, x, ::NoPullbackExtras
+)
+    rc = ruleconfig(backend)
+    y, pullback = rrule_via_ad(rc, f, x)
+    pullbackfunc(dy) = last(pullback(dy))
+    return y, pullbackfunc
+end
+
+function DI.value_and_pullback!_split(
+    f, backend::AutoReverseChainRules, x, extras::NoPullbackExtras
+)
+    y, pullbackfunc = DI.value_and_pullback_split(f, backend, x, extras)
+    pullbackfunc!(dx, dy) = copyto!(dx, pullbackfunc(dy))
+    return y, pullbackfunc!
+end
+
+function DI.value_and_pullback(
+    f, backend::AutoReverseChainRules, x, dy, extras::NoPullbackExtras
+)
+    y, pullbackfunc = DI.value_and_pullback_split(f, backend, x, extras)
+    return y, pullbackfunc(dy)
+end

DifferentiationInterface/src/DifferentiationInterface.jl

@@ -60,6 +60,8 @@ include("sparse/fallbacks.jl")
 include("sparse/jacobian.jl")
 include("sparse/hessian.jl")
 
+include("translation/differentiate_with.jl")
+
 export SecondOrder
 
 export value_and_pushforward!, value_and_pushforward
@@ -87,6 +89,8 @@ export prepare_second_derivative, prepare_hvp, prepare_hessian
 
 export check_available, check_twoarg, check_hessian
 
+export DifferentiateWith
+
 # Re-export backends from ADTypes
 export AutoChainRules
 export AutoDiffractor

DifferentiationInterface/src/translation/differentiate_with.jl

@@ -0,0 +1,54 @@
+"""
+    DifferentiateWith
+
+Callable function wrapper that enforces differentiation with a specified (inner) backend.
+
+This works by defining new rules overriding the behavior of the outer backend that would normally be used.
+
+!!! warning
+    This is an experimental functionality, whose API cannot yet be considered stable.
+    At the moment, it only supports one-argument functions, and rules are only defined for [ChainRules.jl](https://github.com/JuliaDiff/ChainRules.jl)-compatible outer backends.
+
+# Fields
+
+- `f`: the function in question
+- `backend::AbstractADType`: the inner backend to use for differentiation
+
+# Constructor
+
+    DifferentiateWith(f, backend)
+
+# Example
+
+```@repl
+using DifferentiationInterface
+import ForwardDiff, Zygote
+
+function f(x)
+    a = Vector{eltype(x)}(undef, 1)
+    a[1] = sum(x)  # mutation that breaks Zygote
+    return a[1]
+end
+
+dw = DifferentiateWith(f, AutoForwardDiff());
+
+gradient(dw, AutoZygote(), [1.0, 2.0])  # works because it calls ForwardDiff instead
+gradient(f, AutoZygote(), [1.0, 2.0])  # fails
+```
+"""
+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)
+    (; f, backend) = dw
+    return print(io, "$f differentiated with $(backend_str(backend))")
+end

DifferentiationInterface/src/utils/exceptions.jl

@@ -2,7 +2,7 @@ struct MissingBackendError <: Exception
     backend::AbstractADType
 end
 function Base.showerror(io::IO, e::MissingBackendError)
-    println(io, "failed to use $(backend_string(e.backend)) backend.")
+    println(io, "failed to use $(backend_str(e.backend)) backend.")
     if !check_available(e.backend)
         print(
             io,