add AbstractDifferentiation support

Author: oscarddssmith

Project.toml

@@ -4,6 +4,7 @@ authors = ["Keno Fischer <keno@juliacomputing.com> and contributors"]
 version = "0.2.0"
 
 [deps]
+AbstractDifferentiation = "c29ec348-61ec-40c8-8164-b8c60e9d9f3d"
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"

src/AbstractDifferentiation.jl

@@ -0,0 +1,16 @@
+import AbstractDifferentiation as AD
+struct DiffractorForwardBackend <: AD.AbstractForwardMode
+end
+
+bundle(x::Number, dx) = TaylorBundle{1}(x, (dx,))
+bundle(x::Tuple, dx) = CompositeBundle{1}(x, dx)
+bundle(x::AbstractArray{<:Number}, dx::AbstractArray{<:Number}) = TaylorBundle{1}(x, (dx,))  # TODO check me
+# TODO: other types of primal
+
+
+AD.@primitive function pushforward_function(b::DiffractorForwardBackend, f, args...)
+    return function pushforward(vs)
+        z = ∂☆{1}()(ZeroBundle{1}(f), map(bundle, args, vs)...)
+        z[TaylorTangentIndex(1)]
+    end
+end

src/Diffractor.jl

@@ -40,6 +40,7 @@ const GENERATORS = Expr[]
     include("debugutils.jl")
 
     include("stage1/termination.jl")
+    include("AbstractDifferentiation.jl")
 end
 
 end

src/interface.jl

@@ -118,7 +118,7 @@ end
 
 function (f::∇)(args...)
     y, f☆ = ∂⃖(getfield(f, :f), args...)
-    return tail(f☆(dx(y)))
+    tail(f☆(dx(y)))
 end
 
 # N.B: This means the gradient is not available for zero-arg function, but such
@@ -149,9 +149,7 @@ PrimeDerivativeBack(f) = PrimeDerivativeBack{1, typeof(f)}(f)
 PrimeDerivativeBack(f::PrimeDerivativeBack{N, T}) where {N, T} = raise_pd(f)
 
 function (f::PrimeDerivativeBack)(x)
-    z = ∂⃖¹(lower_pd(f), x)
-    y = getfield(z, 1)
-    f☆ = getfield(z, 2)
+    y, f☆ = ∂⃖¹(lower_pd(f), x)
     return unthunk(getfield(f☆(dx(y)), 2))
 end
 
@@ -227,5 +225,3 @@ macro ∂(expr)
 end
 derivative(f, x) = Diffractor.PrimeDerivativeFwd(f)(x)
 const gradient = ∇
-jacobian(f, x::AbstractArray) = reduce(hcat, vec.(gradient(f, x)))
-hessian(f, x::AbstractArray) = jacobian(y -> gradient(f, y), float(x))

test/AbstractDifferentiationTests.jl

@@ -0,0 +1,44 @@
+using AbstractDifferentiation, Diffractor, Test
+include(joinpath(pathof(AbstractDifferentiation), "..", "..", "test", "test_utils.jl"))
+import AbstractDifferentiation as AD
+
+backend = Diffractor.DiffractorForwardBackend()
+@testset "ForwardDiffBackend" begin
+    backends = [
+        @inferred(Diffractor.DiffractorForwardBackend())
+    ]
+    @testset for backend in backends
+        @test backend isa AD.AbstractForwardMode
+
+        @testset "Derivative" begin #setfield!(::Core.Box, ::Symbol, ::Float64)
+            @test_broken test_derivatives(backend)
+        end
+        @testset "Gradient" begin #Diffractor.TangentBundle{1, Float64, Diffractor.TaylorTangent{Tuple{Float64}}}(::Float64, ::Tuple{Float64})
+            @test_broken test_gradients(backend)
+        end
+        @testset "Jacobian" begin #setfield!(::Core.Box, ::Symbol, ::Vector{Float64})
+            @test_broken test_jacobians(backend)
+        end
+        @testset "Hessian" begin #setindex!(::ChainRulesCore.ZeroTangent, ::Float64, ::Int64)
+            @test_broken test_hessians(backend)
+        end
+        @testset "jvp" begin #setfield!(::Core.Box, ::Symbol, ::Vector{Float64})
+            @test_broken test_jvp(backend; vaugmented=true)
+        end
+        @testset "j′vp" begin #setfield!(::Core.Box, ::Symbol, ::Vector{Float64})
+            @test_broken test_j′vp(backend)
+        end
+        @testset "Lazy Derivative" begin
+            test_lazy_derivatives(backend)
+        end
+        @testset "Lazy Gradient" begin #Diffractor.TangentBundle{1, Float64, Diffractor.TaylorTangent{Tuple{Float64}}}(::Float64, ::Tuple{Float64})
+            @test_broken test_lazy_gradients(backend)
+        end
+        @testset "Lazy Jacobian" begin
+            test_lazy_jacobians(backend; vaugmented=true)
+        end
+        @testset "Lazy Hessian" begin # everything everywhere all at once is broken
+            @test_broken test_lazy_hessians(backend)
+        end
+    end
+end

test/runtests.jl

@@ -20,6 +20,7 @@ const bwd = Diffractor.PrimeDerivativeBack
     "forward.jl",
     "reverse.jl",
     "regression.jl",
+    "AbstractDifferentiationTests.jl"
     #"pinn.jl",  # Higher order control flow not yet supported (https://github.com/JuliaDiff/Diffractor.jl/issues/24)
 )
     include(file)