Jacobian and friends

Project.toml

@@ -5,6 +5,7 @@ version = "0.1.0"
 
 [deps]
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [weakdeps]
@@ -14,27 +15,25 @@ Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
 DifferentiationInterfaceChainRulesCoreExt = "ChainRulesCore"
 DifferentiationInterfaceEnzymeExt = "Enzyme"
 DifferentiationInterfaceFiniteDiffExt = "FiniteDiff"
 DifferentiationInterfaceForwardDiffExt = ["ForwardDiff", "DiffResults"]
 DifferentiationInterfaceReverseDiffExt = ["ReverseDiff", "DiffResults"]
+DifferentiationInterfaceZygoteExt = ["Zygote"]
 
 [compat]
 ChainRulesCore = "1.19"
 DiffResults = "1.1"
 DocStringExtensions = "0.9"
-FiniteDiff = "2.22"
 Enzyme = "0.11"
+FillArrays = "1"
+FiniteDiff = "2.22"
 ForwardDiff = "0.10"
 LinearAlgebra = "1"
 ReverseDiff = "1.15"
+Zygote = "0.6"
 julia = "1.10"
-
-[extras]
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-
-[targets]
-test = ["Test"]

docs/Project.toml

@@ -1,6 +1,13 @@
 [deps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
+Diffractor = "9f5e2b26-1114-432f-b630-d3fe2085c51c"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 Documenter = "1"

docs/make.jl

@@ -1,6 +1,21 @@
 using DifferentiationInterface
+import DifferentiationInterface as DI
 using Documenter
 
+using Diffractor: Diffractor
+using Enzyme: Enzyme
+using FiniteDiff: FiniteDiff
+using ForwardDiff: ForwardDiff
+using ReverseDiff: ReverseDiff
+using Zygote: Zygote
+
+DIChainRulesCoreExt = Base.get_extension(DI, :DifferentiationInterfaceChainRulesCoreExt)
+DIEnzymeExt = Base.get_extension(DI, :DifferentiationInterfaceEnzymeExt)
+DIFiniteDiffExt = Base.get_extension(DI, :DifferentiationInterfaceFiniteDiffExt)
+DIForwardDiffExt = Base.get_extension(DI, :DifferentiationInterfaceForwardDiffExt)
+DIReverseDiffExt = Base.get_extension(DI, :DifferentiationInterfaceReverseDiffExt)
+DIZygoteExt = Base.get_extension(DI, :DifferentiationInterfaceZygoteExt)
+
 DocMeta.setdocmeta!(
     DifferentiationInterface,
     :DocTestSetup,
@@ -9,16 +24,24 @@ DocMeta.setdocmeta!(
 )
 
 makedocs(;
-    modules=[DifferentiationInterface],
+    modules=[
+        DifferentiationInterface,
+        DIChainRulesCoreExt,
+        DIEnzymeExt,
+        DIFiniteDiffExt,
+        DIForwardDiffExt,
+        DIReverseDiffExt,
+        DIZygoteExt,
+    ],
     authors="Guillaume Dalle, Adrian Hill",
     sitename="DifferentiationInterface.jl",
     format=Documenter.HTML(;
         prettyurls=get(ENV, "CI", "false") == "true",
         canonical="https://gdalle.github.io/DifferentiationInterface.jl",
-        edit_link="main",
-        assets=String[],
     ),
-    pages=["Home" => "index.md"],
+    pages=[
+        "Home" => "index.md", "API reference" => "api.md", "Extensions" => "extensions.md"
+    ],
 )
 
 deploydocs(; repo="github.com/gdalle/DifferentiationInterface.jl", devbranch="main")

docs/src/api.md

@@ -0,0 +1,70 @@
+```@meta
+CurrentModule = DifferentiationInterface
+CollapsedDocStrings = true
+```
+
+# API reference
+
+```@docs
+DifferentiationInterface
+```
+
+## Backends
+
+```@autodocs
+Modules = [DifferentiationInterface]
+Pages = ["backends.jl"]
+```
+
+## Primitives
+
+### Pushforward
+
+```@autodocs
+Modules = [DifferentiationInterface]
+Pages = ["pushforward.jl"]
+```
+
+### Pullback
+
+```@autodocs
+Modules = [DifferentiationInterface]
+Pages = ["pullback.jl"]
+```
+
+## Special cases
+
+### Scalar to scalar
+
+```@autodocs
+Modules = [DifferentiationInterface]
+Pages = ["scalar_scalar.jl"]
+```
+
+### Scalar to array
+
+```@autodocs
+Modules = [DifferentiationInterface]
+Pages = ["scalar_array.jl"]
+```
+
+### Array to scalar
+
+```@autodocs
+Modules = [DifferentiationInterface]
+Pages = ["array_scalar.jl"]
+```
+
+### Array to array
+
+```@autodocs
+Modules = [DifferentiationInterface]
+Pages = ["array_array.jl"]
+```
+
+## Internals
+
+```@autodocs
+Modules = [DifferentiationInterface]
+Public = false
+```

docs/src/extensions.md

@@ -0,0 +1,41 @@
+```@meta
+CollapsedDocStrings = true
+```
+
+# Extensions
+
+## ChainRulesCore
+
+```@autodocs
+Modules = [DIChainRulesCoreExt]
+```
+
+## Enzyme
+
+```@autodocs
+Modules = [DIEnzymeExt]
+```
+
+## FiniteDiff
+
+```@autodocs
+Modules = [DIFiniteDiffExt]
+```
+
+## ForwardDiff
+
+```@autodocs
+Modules = [DIForwardDiffExt]
+```
+
+## ReverseDiff
+
+```@autodocs
+Modules = [DIReverseDiffExt]
+```
+
+## Zygote
+
+```@autodocs
+Modules = [DIZygoteExt]
+```

docs/src/index.md

@@ -1,39 +1,17 @@
 ```@meta
 CurrentModule = DifferentiationInterface
+CollapsedDocStrings = true
 ```
 
 # DifferentiationInterface
 
 Documentation for [DifferentiationInterface](https://github.com/gdalle/DifferentiationInterface.jl).
 
-This is an interface to various autodiff backends for differentiating functions of the form `f(x) = y`, where `x` and `y` are either numbers or arrays.
+This is an interface to various autodiff backends for differentiating functions of the form `f(x) = y`, where `x` and `y` are either numbers or abstract arrays.
 
-## Public
+## Terminology
 
-```@autodocs
-Modules = [DifferentiationInterface]
-Private = false
-```
-
-## Internals
-
-```@autodocs
-Modules = [DifferentiationInterface]
-Public = false
-```
-
-## Math
-
-Some implementation reminders:
-
-|                  | pushforward                                                      | pullback                                                           |
-| ---------------- | ---------------------------------------------------------------- | ------------------------------------------------------------------ |
-| scalar -> scalar | derivative multiplied by the tangent                             | derivative multiplied by the cotangent                             |
-| scalar -> vector | derivative vector multiplied componentwise by the tangent vector | dot product between the derivative vector and the cotangent vector |
-| vector -> scalar | dot product between the gradient vector and the tangent vector   | gradient vector multiplied componentwise by the cotangent          |
-| vector -> vector | Jacobian matrix multiplied by the tangent vector                 | transposed Jacobian matrix multiplied by the cotangent vector      |
-
-## Index
-
-```@index
-```
+|              | scalar output | vector output   |
+| ------------ | ------------- | --------------- |
+| scalar input | derivative    | multiderivative |
+| vector input | gradient      | jacobian        |

ext/DifferentiationInterfaceChainRulesCoreExt.jl

@@ -2,23 +2,30 @@ module DifferentiationInterfaceChainRulesCoreExt
 
 using ChainRulesCore: NoTangent, frule_via_ad, rrule_via_ad
 using DifferentiationInterface
+using DocStringExtensions
 
 ruleconfig(backend::ChainRulesForwardBackend) = backend.ruleconfig
 ruleconfig(backend::ChainRulesReverseBackend) = backend.ruleconfig
 
 update!(_old::Number, new::Number) = new
 update!(old, new) = old .= new
 
+"""
+$(TYPEDSIGNATURES)
+"""
 function DifferentiationInterface.value_and_pushforward!(
-    dy::Y, backend::ChainRulesForwardBackend, f, x::X, dx::X
+    dy::Y, backend::ChainRulesForwardBackend, f, x::X, dx
 ) where {X,Y}
     rc = ruleconfig(backend)
     y, new_dy = frule_via_ad(rc, (NoTangent(), dx), f, x)
     return y, update!(dy, new_dy)
 end
 
+"""
+$(TYPEDSIGNATURES)
+"""
 function DifferentiationInterface.value_and_pullback!(
-    dx::X, backend::ChainRulesReverseBackend, f, x::X, dy::Y
+    dx, backend::ChainRulesReverseBackend, f, x::X, dy::Y
 ) where {X,Y}
     rc = ruleconfig(backend)
     y, pullback = rrule_via_ad(rc, f, x)

ext/DifferentiationInterfaceEnzymeExt.jl

@@ -1,19 +1,39 @@
 module DifferentiationInterfaceEnzymeExt
 
 using DifferentiationInterface
+using DocStringExtensions
 using Enzyme: Forward, ReverseWithPrimal, Active, Duplicated, autodiff
 
+const EnzymeBackends = Union{EnzymeForwardBackend,EnzymeReverseBackend}
+
+## Unit vector
+
+# Enzyme's `Duplicated(x, dx)` expects both arguments to be of the same type
+function DifferentiationInterface.basisarray(
+    ::EnzymeBackends, a::AbstractArray{T}, i
+) where {T}
+    b = zero(a)
+    b[i] = one(T)
+    return b
+end
+
 ## Forward mode
 
+"""
+$(TYPEDSIGNATURES)
+"""
 function DifferentiationInterface.value_and_pushforward!(
-    _dy::Y, ::EnzymeForwardBackend, f, x::X, dx::X
+    _dy::Y, ::EnzymeForwardBackend, f, x::X, dx
 ) where {X,Y<:Real}
     y, new_dy = autodiff(Forward, f, Duplicated, Duplicated(x, dx))
     return y, new_dy
 end
 
+"""
+$(TYPEDSIGNATURES)
+"""
 function DifferentiationInterface.value_and_pushforward!(
-    dy::Y, ::EnzymeForwardBackend, f, x::X, dx::X
+    dy::Y, ::EnzymeForwardBackend, f, x::X, dx
 ) where {X,Y<:AbstractArray}
     y, new_dy = autodiff(Forward, f, Duplicated, Duplicated(x, dx))
     dy .= new_dy
@@ -23,7 +43,7 @@ end
 ## Reverse mode
 
 function DifferentiationInterface.value_and_pullback!(
-    _dx::X, ::EnzymeReverseBackend, f, x::X, dy::Y
+    _dx, ::EnzymeReverseBackend, f, x::X, dy::Y
 ) where {X<:Number,Y<:Union{Real,Nothing}}
     der, y = autodiff(ReverseWithPrimal, f, Active, Active(x))
     new_dx = dy * only(der)