Merge pull request #187 from oscardssmith/add-interface-v2

Give Diffractor an AbstractDifferentiation interface

Author: oxinabox

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"
@@ -14,6 +15,13 @@ PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 
+[extras]
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
 [compat]
 ChainRules = "1.44.6"
 ChainRulesCore = "1.15.3"
@@ -24,3 +32,6 @@ PrecompileTools = "1"
 StaticArrays = "1"
 StructArrays = "0.6"
 julia = "1.10"
+
+[targets]
+test = ["ForwardDiff", "LinearAlgebra", "Random", "Symbolics", "Test"]

docs/src/index.md

@@ -3,3 +3,32 @@
 Next-generation AD
 
 [PDF containing the terminology](terminology.pdf)
+
+## Getting Started
+
+⚠️This certainly has bugs and issues. Please open issues on [Diffractor.jl](https://github.com/JuliaDiff/Diffractor.jl/), or [AbstractDifferentiation.jl](https://github.com/JuliaDiff/AbstractDifferentiation.jl/issues) as appropriate.⚠️
+
+Diffractor's public API is via [AbstractDifferentiation.jl](https://github.com/JuliaDiff/AbstractDifferentiation.jl/).
+Please see the AbstractDifferentiation.jl docs for detailed usage.
+
+```@jldoctest
+julia> using Diffractor: DiffractorForwardBackend
+
+julia> using AbstractDifferentiation: derivative
+
+julia> derivative(DiffractorForwardBackend(), +, 1.5, 10.0)
+(1.0, 1.0)
+
+julia> derivative(DiffractorForwardBackend(), *, 1.5, 10.0)
+(10.0, 1.5)
+
+julia> jacobian(DiffractorForwardBackend(), prod, [1.5, 2.5, 10.0]) |> only
+1×3 Matrix{Float64}:
+ 25.0  15.0  3.75
+
+julia> jacobian(DiffractorForwardBackend(), identity, [1.5, 2.5, 10.0]) |> only
+3×3 Matrix{Float64}:
+ 1.0  0.0  0.0
+ 0.0  1.0  0.0
+ 0.0  0.0  1.0
+```

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
@@ -127,8 +127,6 @@ function (::Type{∇})(f, x1, args...)
     unthunk.(∇(f)(x1, args...))
 end
 
-const gradient = ∇
-
 # Star Trek has their prime directive. We have the...
 abstract type AbstractPrimeDerivative{N, T}; end
 
@@ -151,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
 
@@ -181,8 +177,8 @@ struct PrimeDerivative{N, T}
 end
 
 function (f::PrimeDerivative{N, T})(x) where {N, T}
-    # For now, this is backwards mode, since that's more fully implemented
-    return PrimeDerivativeBack{N, T}(f.f)(x)
+    # For now, this is forward mode, since that's more fully implemented
+    return PrimeDerivativeFwd{N, T}(f.f)(x)
 end
 
 """
@@ -227,3 +223,5 @@ will compute the derivative `∂^3 f/∂x^2 ∂y` at `(x,y)`.
 macro ∂(expr)
     error("Write me")
 end
+derivative(f, x) = Diffractor.PrimeDerivativeFwd(f)(x)
+const gradient = ∇

test/AbstractDifferentiationTests.jl

@@ -0,0 +1,52 @@
+using AbstractDifferentiation, Diffractor, Test, LinearAlgebra
+import AbstractDifferentiation as AD
+backend = Diffractor.DiffractorForwardBackend()
+
+@testset "basics" begin
+    @test AD.derivative(backend, +, 1.5, 10.0) == (1.0, 1.0)
+    @test AD.derivative(backend, *, 1.5, 10.0) == (10.0, 1.5)
+    @test only(AD.jacobian(backend, prod, [1.5, 2.5, 10.0])) == [25.0  15.0  3.75]
+    @test only(AD.jacobian(backend, identity, [1.5, 2.5, 10.0])) == Matrix(I, 3, 3)
+end
+
+# standard tests from AbstractDifferentiation.test_utils
+include(joinpath(pathof(AbstractDifferentiation), "..", "..", "test", "test_utils.jl"))
+@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 #MethodError: no method matching *(::Diffractor.PrimeDerivativeBack{1, Diagonal{Bool, Vector{Bool}}}, ::Vector{Float64})
+            @test_broken 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/Project.toml

@@ -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)