Merge pull request #205 from JuliaDiff/ox/bundle_helper

add more types supported with bundle helper

Author: oxinabox

src/AbstractDifferentiation.jl

@@ -2,10 +2,27 @@ 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
+"""
+    bundle(primal, tangent)
+
+Wraps a primal up with a tangent into the appropriate kind of `AbstractBundle{1}`.
+This is more or less the Diffractor equivelent of ForwardDiff.jl's `Dual` type.
+"""
+function bundle end
+bundle(x, dx::ChainRulesCore.AbstractZero) = UniformBundle{1, typeof(x), typeof(dx)}(x, dx)
+bundle(x::Number, dx::Number) = TaylorBundle{1}(x, (dx,))
+bundle(x::AbstractArray{<:Number}, dx) = TaylorBundle{1}(x, (dx,))
+bundle(x::AbstractArray, dx) = error("Nonnumeric arrays not implemented, that type is a mess")
+bundle(x::P, dx::Tangent{P}) where P = _bundle(x, ChainRulesCore.canonicalize(dx))
+
+"helper that assumes tangent is in canonical form"
+function _bundle(x::P, dx::Tangent{P}) where P
+    # SoA to AoS flip (hate this, hate it even more cos we just undo it later when we hit chainrules)
+    the_bundle = ntuple(Val{fieldcount(P)}()) do ii
+        bundle(getfield(x, ii), getproperty(dx, ii))
+    end
+    return CompositeBundle{1, P}(the_bundle)
+end
 
 
 AD.@primitive function pushforward_function(b::DiffractorForwardBackend, f, args...)

test/AbstractDifferentiationTests.jl

@@ -1,7 +1,25 @@
-using AbstractDifferentiation, Diffractor, Test, LinearAlgebra
+using AbstractDifferentiation, Diffractor, Test, LinearAlgebra, ChainRulesCore
 import AbstractDifferentiation as AD
 backend = Diffractor.DiffractorForwardBackend()
 
+@testset "bundle" begin
+    bundle = Diffractor.bundle
+
+    @test bundle(1.0, 2.0) isa Diffractor.TaylorBundle{1}
+    @test bundle([1.0, 2.0], [2.0, 3.0]) isa Diffractor.TaylorBundle{1}
+    @test bundle(1.5=>2.5, Tangent{Pair{Float64, Float64}}(first=1.0, second=2.0)) isa Diffractor.CompositeBundle{1}
+    @test bundle(1.1, ChainRulesCore.ZeroTangent()) isa Diffractor.ZeroBundle{1}
+    @test bundle(1.5=>2.5=>3.5, Tangent{Pair{Float64, Pair{Float64, Float64}}}(first=1.0, second=Tangent{Pair{Float64, Float64}}(first=1.0, second=2.0))) isa Diffractor.CompositeBundle{1}
+    
+    # noncanonical structural tangent
+    b = bundle(1.5=>2.5=>3.5, Tangent{Pair{Float64, Pair{Float64, Float64}}}(second=Tangent{Pair{Float64, Float64}}(second=2.0, first=1.0)))
+    t = Diffractor.first_partial(b)
+    @test b isa Diffractor.CompositeBundle{1}
+    @test iszero(t.first)
+    @test t.second.first == 1.0
+    @test t.second.second == 2.0
+end
+
 @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)
@@ -50,3 +68,4 @@ include(joinpath(pathof(AbstractDifferentiation), "..", "..", "test", "test_util
         end
     end
 end
+