Gradient for AutoFastDifferentiation + some renaming (#180)

* Gradient for AutoFastDifferentiation + some renaming

* Reactivate first order

* Documenter mutation better

Author: gdalle

DifferentiationInterface/README.md

@@ -30,19 +30,19 @@ This package provides a backend-agnostic syntax to differentiate functions of th
 
 We support most of the backends defined by [ADTypes.jl](https://github.com/SciML/ADTypes.jl):
 
-| Backend                                                                         | Object                                                     |
-| :------------------------------------------------------------------------------ | :--------------------------------------------------------- |
-| [ChainRulesCore.jl](https://github.com/JuliaDiff/ChainRulesCore.jl)             | `AutoChainRules(ruleconfig)`                               |
-| [Diffractor.jl](https://github.com/JuliaDiff/Diffractor.jl)                     | `AutoDiffractor()`                                         |
-| [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl)                              | `AutoEnzyme(Enzyme.Forward)`, `AutoEnzyme(Enzyme.Reverse)` |
-| [FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl)                     | `AutoFiniteDiff()`                                         |
-| [FiniteDifferences.jl](https://github.com/JuliaDiff/FiniteDifferences.jl)       | `AutoFiniteDifferences(fdm)`                               |
-| [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl)                   | `AutoForwardDiff()`                                        |
-| [PolyesterForwardDiff.jl](https://github.com/JuliaDiff/PolyesterForwardDiff.jl) | `AutoPolyesterForwardDiff(; chunksize)`                    |
-| [ReverseDiff.jl](https://github.com/JuliaDiff/ReverseDiff.jl)                   | `AutoReverseDiff()`                                        |
-| [SparseDiffTools.jl](https://github.com/JuliaDiff/SparseDiffTools.jl)           | `AutoSparseForwardDiff()`, `AutoSparseFiniteDiff()`        |
-| [Tracker.jl](https://github.com/FluxML/Tracker.jl)                              | `AutoTracker()`                                            |
-| [Zygote.jl](https://github.com/FluxML/Zygote.jl)                                | `AutoZygote()`                                             |
+| Backend                                                                         | Object                                                                   |
+| :------------------------------------------------------------------------------ | :----------------------------------------------------------------------- |
+| [ChainRulesCore.jl](https://github.com/JuliaDiff/ChainRulesCore.jl)             | `AutoChainRules(; ruleconfig)`                                           |
+| [Diffractor.jl](https://github.com/JuliaDiff/Diffractor.jl)                     | `AutoDiffractor()`                                                       |
+| [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl)                              | `AutoEnzyme(; mode=Enzyme.Forward)`, `AutoEnzyme(; mode=Enzyme.Reverse)` |
+| [FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl)                     | `AutoFiniteDiff()`                                                       |
+| [FiniteDifferences.jl](https://github.com/JuliaDiff/FiniteDifferences.jl)       | `AutoFiniteDifferences(; fdm)`                                           |
+| [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl)                   | `AutoForwardDiff()`                                                      |
+| [PolyesterForwardDiff.jl](https://github.com/JuliaDiff/PolyesterForwardDiff.jl) | `AutoPolyesterForwardDiff(; chunksize)`                                  |
+| [ReverseDiff.jl](https://github.com/JuliaDiff/ReverseDiff.jl)                   | `AutoReverseDiff()`                                                      |
+| [SparseDiffTools.jl](https://github.com/JuliaDiff/SparseDiffTools.jl)           | `AutoSparseForwardDiff()`, `AutoSparseFiniteDiff()`                      |
+| [Tracker.jl](https://github.com/FluxML/Tracker.jl)                              | `AutoTracker()`                                                          |
+| [Zygote.jl](https://github.com/FluxML/Zygote.jl)                                | `AutoZygote()`                                                           |
 
 We also provide some experimental backends ourselves:
 

DifferentiationInterface/docs/src/api.md

@@ -87,7 +87,7 @@ value_and_pullback!_split
 
 ```@docs
 check_available
-check_mutation
+check_twoarg
 check_hessian
 ```
 

DifferentiationInterface/docs/src/backends.md

@@ -93,13 +93,13 @@ Markdown.parse(join(vcat(header, subheader, rows...), "\n"))  # hide
 
 All backends are compatible with one-argument functions `f(x) = y`.
 Only some are compatible with two-argument functions `f!(y, x) = nothing`.
-You can use [`check_mutation`](@ref) to check that feature, like we did below:
+You can use [`check_twoarg`](@ref) to check that feature, like we did below:
 
 ```@example backends
 header = "| backend | mutation |"  # hide
 subheader = "|:---|:---:|"  # hide
 rows = map(all_backends()) do backend  # hide
-    "| `$(backend_string(backend))` | $(check_mutation(backend) ? '✅' : '❌') |"  # hide
+    "| `$(backend_string(backend))` | $(check_twoarg(backend) ? '✅' : '❌') |"  # hide
 end  # hide
 Markdown.parse(join(vcat(header, subheader, rows...), "\n"))  # hide
 ```

DifferentiationInterface/docs/src/overview.md

@@ -47,13 +47,16 @@ Several variants of each operator are defined:
 In order to ensure symmetry between one-argument functions `f(x) = y` and two-argument functions `f!(y, x) = nothing`, we define the same operators for both cases.
 However they have different signatures:
 
-| signature  | out-of-place                       | in-place                                 |
-| :--------- | :--------------------------------- | :--------------------------------------- |
-| `f(x)`     | `operator(f,     backend, x, ...)` | `operator!(f,     res, backend, x, ...)` |
-| `f!(y, x)` | `operator(f!, y, backend, x, ...)` | `operator!(f!, y, res, backend, x, ...)` |
+| signature  | out-of-place                       | in-place                                    |
+| :--------- | :--------------------------------- | :------------------------------------------ |
+| `f(x)`     | `operator(f,     backend, x, ...)` | `operator!(f,     result, backend, x, ...)` |
+| `f!(y, x)` | `operator(f!, y, backend, x, ...)` | `operator!(f!, y, result, backend, x, ...)` |
 
 !!! warning
-    Every variant of the operator will mutate `y` when applied to a two-argument function `f!(y, x) = nothing`, even if it does not have a `!` in its name.
+    Our mutation convention is that all positional arguments between `f`/`f!` and `backend` are mutated (the `extras` as well, see below).
+    This convention holds regardless of the bang `!` in the operator name, because we assume that a user passing a two-argument function `f!(y, x)` anticipates mutation anyway.
+
+    Still, better be careful with two-argument functions, because every variant of the operator will mutate `y`... even if it does not have a `!` in its name (see the bottom left cell in the table).
 
 ## Preparation
 
@@ -71,14 +74,21 @@ This is a backend-specific procedure, but we expose a common syntax to achieve i
 | `pullback`          | [`prepare_pullback`](@ref)          |
 | `hvp`               | [`prepare_hvp`](@ref)               |
 
-If you run `prepare_operator(backend, f, x, [seed])`, it will create an object called `extras` containing the necessary information to speed up `operator` and its variants.
+Unsurprisingly, preparation syntax depends on the number of arguments:
+
+| signature  | preparation signature                      |
+| :--------- | :----------------------------------------- |
+| `f(x)`     | `prepare_operator(f,     backend, x, ...)` |
+| `f!(y, x)` | `prepare_operator(f!, y, backend, x, ...)` |
+
+The preparation `prepare_operator(f, backend, x)` will create an object called `extras` containing the necessary information to speed up `operator` and its variants.
 This information is specific to `backend` and `f`, as well as the _type and size_ of the input `x` and the _control flow_ within the function, but it should work with different _values_ of `x`.
 
-You can then call `operator(backend, f, x2, extras)`, which should be faster than `operator(f, backend, x2)`.
+You can then call e.g. `operator(backend, f, x2, extras)`, which should be faster than `operator(f, backend, x2)`.
 This is especially worth it if you plan to call `operator` several times in similar settings: you can think of it as a warm up.
 
 !!! warning
-    The `extras` object is nearly always mutated, even if the operator does not have a `!` in its name.
+    The `extras` object is nearly always mutated when given to an operator, even when said operator does not have a bang `!` in its name.
 
 ### Second order
 
@@ -129,7 +139,7 @@ We make this available for all backends with the following operators:
 ### Non-standard types
 
 The package is thoroughly tested with inputs and outputs of the following types: `Float64`, `Vector{Float64}` and `Matrix{Float64}`.
-We also expect it to work on all kinds of `Number` and `AbstractArray` variables.
+We also expect it to work on most kinds of `Number` and `AbstractArray` variables.
 Beyond that, you are in uncharted territory.
 We voluntarily keep the type annotations minimal, so that passing more complex objects or custom structs _might work with some backends_, but we make no guarantees about that.
 

DifferentiationInterface/docs/src/tutorial.md

@@ -63,10 +63,12 @@ Some backends get a speed boost from this trick.
 
 ```@repl tuto
 grad = zero(x)
-grad = gradient!(f, grad, backend, x)
+gradient!(f, grad, backend, x);
+grad
 ```
 
-Note the double exclamation mark, which is a convention telling you that `grad` _may or may not_ be overwritten, but will be returned either way (see [this section](@ref Variants) for more details).
+The bang indicates that one of the arguments of `gradient!` might be mutated.
+More precisely, our convention is that _every positional argument between the function and the backend is mutated (and the `extras` too, see below)_.
 
 ```@repl tuto
 @btime gradient!($f, _grad, $backend, $x) evals=1 setup=(_grad=similar($x));
@@ -90,7 +92,8 @@ You don't need to know what this object is, you just need to pass it to the grad
 
 ```@repl tuto
 grad = zero(x);
-grad = gradient!(f, grad, backend, x, extras)
+gradient!(f, grad, backend, x, extras);
+grad
 ```
 
 Preparation makes the gradient computation much faster, and (in this case) allocation-free.
@@ -102,6 +105,8 @@ Preparation makes the gradient computation much faster, and (in this case) alloc
 );
 ```
 
+Beware that the `extras` object is nearly always mutated by differentiation operators, even though it is given as the last positional argument.
+
 ## Switching backends
 
 The whole point of DifferentiationInterface.jl is that you can easily experiment with different AD solutions.

DifferentiationInterface/ext/DifferentiationInterfaceChainRulesCoreExt/DifferentiationInterfaceChainRulesCoreExt.jl

@@ -12,7 +12,7 @@ const AutoForwardChainRules = AutoChainRules{<:RuleConfig{>:HasForwardsMode}}
 const AutoReverseChainRules = AutoChainRules{<:RuleConfig{>:HasReverseMode}}
 
 DI.check_available(::AutoChainRules) = true
-DI.supports_mutation(::AutoChainRules) = DI.MutationNotSupported()
+DI.mutation_support(::AutoChainRules) = DI.MutationNotSupported()
 DI.mode(::AutoForwardChainRules) = ADTypes.AbstractForwardMode
 DI.mode(::AutoReverseChainRules) = ADTypes.AbstractReverseMode
 

DifferentiationInterface/ext/DifferentiationInterfaceDiffractorExt/DifferentiationInterfaceDiffractorExt.jl

@@ -6,7 +6,7 @@ using DifferentiationInterface: NoPushforwardExtras
 using Diffractor: DiffractorRuleConfig, TaylorTangentIndex, ZeroBundle, bundle, ∂☆
 
 DI.check_available(::AutoDiffractor) = true
-DI.supports_mutation(::AutoDiffractor) = DI.MutationNotSupported()
+DI.mutation_support(::AutoDiffractor) = DI.MutationNotSupported()
 DI.mode(::AutoDiffractor) = ADTypes.AbstractForwardMode
 DI.mode(::AutoChainRules{<:DiffractorRuleConfig}) = ADTypes.AbstractForwardMode
 

DifferentiationInterface/ext/DifferentiationInterfaceFastDifferentiationExt/DifferentiationInterfaceFastDifferentiationExt.jl

@@ -31,7 +31,7 @@ const AnyAutoFastDifferentiation = Union{
     AutoFastDifferentiation,AutoSparseFastDifferentiation
 }
 
-DI.check_available(::AutoFastDifferentiation) = true
+DI.check_available(::AnyAutoFastDifferentiation) = true
 DI.mode(::AnyAutoFastDifferentiation) = ADTypes.AbstractSymbolicDifferentiationMode
 DI.pushforward_performance(::AnyAutoFastDifferentiation) = DI.PushforwardFast()
 DI.pullback_performance(::AnyAutoFastDifferentiation) = DI.PullbackSlow()

DifferentiationInterface/ext/DifferentiationInterfaceFastDifferentiationExt/onearg.jl

@@ -136,6 +136,64 @@ function DI.value_and_derivative!(
     return f(x), DI.derivative!(f, der, backend, x, extras)
 end
 
+## Gradient
+
+struct FastDifferentiationOneArgGradientExtras{E1,E2} <: GradientExtras
+    jac_exe::E1
+    jac_exe!::E2
+end
+
+function DI.prepare_gradient(f, backend::AnyAutoFastDifferentiation, x)
+    y_prototype = f(x)
+    x_var = make_variables(:x, size(x)...)
+    y_var = f(x_var)
+
+    x_vec_var = vec(x_var)
+    y_vec_var = monovec(y_var)
+    jac_var = jacobian(y_vec_var, x_vec_var)
+    jac_exe = make_function(jac_var, x_vec_var; in_place=false)
+    jac_exe! = make_function(jac_var, x_vec_var; in_place=true)
+    return FastDifferentiationOneArgGradientExtras(jac_exe, jac_exe!)
+end
+
+function DI.gradient(
+    f, ::AnyAutoFastDifferentiation, x, extras::FastDifferentiationOneArgGradientExtras
+)
+    jac = extras.jac_exe(vec(x))
+    grad_vec = @view jac[1, :]
+    return reshape(grad_vec, size(x))
+end
+
+function DI.gradient!(
+    f,
+    grad,
+    ::AnyAutoFastDifferentiation,
+    x,
+    extras::FastDifferentiationOneArgGradientExtras,
+)
+    extras.jac_exe!(reshape(grad, 1, length(grad)), vec(x))
+    return grad
+end
+
+function DI.value_and_gradient(
+    f,
+    backend::AnyAutoFastDifferentiation,
+    x,
+    extras::FastDifferentiationOneArgGradientExtras,
+)
+    return f(x), DI.gradient(f, backend, x, extras)
+end
+
+function DI.value_and_gradient!(
+    f,
+    grad,
+    backend::AnyAutoFastDifferentiation,
+    x,
+    extras::FastDifferentiationOneArgGradientExtras,
+)
+    return f(x), DI.gradient!(f, grad, backend, x, extras)
+end
+
 ## Jacobian
 
 struct FastDifferentiationOneArgJacobianExtras{Y,E1,E2} <: JacobianExtras

DifferentiationInterface/ext/DifferentiationInterfaceFiniteDifferencesExt/DifferentiationInterfaceFiniteDifferencesExt.jl

@@ -9,7 +9,7 @@ using FiniteDifferences: FiniteDifferences, grad, jacobian, jvp, j′vp
 using LinearAlgebra: dot
 
 DI.check_available(::AutoFiniteDifferences) = true
-DI.supports_mutation(::AutoFiniteDifferences) = DI.MutationNotSupported()
+DI.mutation_support(::AutoFiniteDifferences) = DI.MutationNotSupported()
 
 function FiniteDifferences.to_vec(a::OneElement)  # TODO: remove type piracy (https://github.com/JuliaDiff/FiniteDifferences.jl/issues/141)
     return FiniteDifferences.to_vec(collect(a))