Initial pass at switching CompositeBundle over to TaylorTangentBundle

Author: oxinabox

src/stage1/forward.jl

@@ -51,6 +51,18 @@ function shuffle_down(b::CompositeBundle{N, B}) where {N, B}
     z
 end
 
+function shuffle_up(r::TaylorBundle{1, Tuple{B1,B2}}) where {B1,B2}
+    z₀ = primal(r)[1]
+    z₁ = partial(r, 1)[1]
+    z₂ = primal(r)[2]
+    z₁₂ = partial(r, 1)[2]
+    if z₁ == z₂
+        return TaylorBundle{2}(z₀, (z₁, z₁₂))
+    else
+        return ExplicitTangentBundle{2}(z₀, (z₁, z₂, z₁₂))
+    end
+end
+
 function shuffle_up(r::CompositeBundle{1})
     z₀ = primal(r.tup[1])
     z₁ = partial(r.tup[1], 1)
@@ -76,6 +88,7 @@ isswifty(::UniformBundle) = true
 isswifty(b::CompositeBundle) = all(isswifty, b.tup)
 isswifty(::Any) = false
 
+#TODO: port this to TaylorTangent:
 function shuffle_up(r::CompositeBundle{N}) where {N}
     a, b = r.tup
     if isswifty(a) && isswifty(b) && taylor_compatible(a, b)

src/stage1/recurse_fwd.jl

@@ -4,13 +4,47 @@ struct ∂vararg{N}; end
 
 (::∂vararg{N})() where {N} = ZeroBundle{N}(())
 function (::∂vararg{N})(a::AbstractTangentBundle{N}...) where N
-    CompositeBundle{N, Tuple{map(x->basespace(typeof(x)), a)...}}(a)
+    B = Tuple{map(x->basespace(Core.Typeof(x)), a)...}
+    return (∂☆new{N}())(B, a...)
 end
 
 struct ∂☆new{N}; end
 
-(::∂☆new{N})(B::Type, a::AbstractTangentBundle{N}...) where {N} =
-    CompositeBundle{N, B}(a)
+# we split out the 1st order derivative as a special case for performance
+# but the nth order case does also work for this
+function (::∂☆new{1})(B::Type, xs::AbstractTangentBundle{1}...)
+    primal_args = map(primal, xs)
+    the_primal = B <: Tuple ? B(primal_args) : B(primal_args...)
+        
+    tangent_tup = map(x->partial(x, 1), xs)
+    the_partial = if B<:Tuple
+        Tangent{B, typeof(tangent_tup)}(tangent_tup)
+    else
+        names = fieldnames(B)
+        tangent_nt = NamedTuple{names}(tangent_tup)
+        Tangent{B, typeof(tangent_nt)}(tangent_nt)
+    end
+    return TaylorBundle{1, B}(the_primal, (the_partial,))
+end
+
+function (::∂☆new{N})(B::Type, xs::AbstractTangentBundle{N}...) where {N}
+    primal_args = map(primal, xs)
+    the_primal = B <: Tuple ? B(primal_args) : B(primal_args...)
+        
+    the_partials = ntuple(Val{N}()) do ii
+        iith_order_type = ii==1 ? B : Any  # the type of the higher order tangents isn't worth tracking
+        tangent_tup = map(x->partial(x, ii), xs)
+        tangent = if B<:Tuple
+            Tangent{iith_order_type, typeof(tangent_tup)}(tangent_tup)
+        else
+            names = fieldnames(B)
+            tangent_nt = NamedTuple{names}(tangent_tup)
+            Tangent{iith_order_type, typeof(tangent_nt)}(tangent_nt)
+        end
+        return tangent
+    end
+    return TaylorBundle{N, B}(the_primal, the_partials)
+end
 
 @generated (::∂☆new{N})(B::Type) where {N} = return :(ZeroBundle{$N}($(Expr(:new, :B))))
 

test/forward.jl

@@ -25,8 +25,7 @@ let var"'" = Diffractor.PrimeDerivativeFwd
     # Integration tests
     @test recursive_sin'(1.0) == cos(1.0)
     @test recursive_sin''(1.0) == -sin(1.0)
-    # Error: ArgumentError: Tangent for the primal Tangent{Tuple{Float64, Float64}, Tuple{Float64, Float64}}
-    # should be backed by a NamedTuple type, not by Tuple{Tangent{Tuple{Float64, Float64}, Tuple{Float64, Float64}}}.
+    
     @test_broken recursive_sin'''(1.0) == -cos(1.0)
     @test_broken recursive_sin''''(1.0) == sin(1.0)
     @test_broken recursive_sin'''''(1.0) == cos(1.0)
@@ -40,6 +39,7 @@ let var"'" = Diffractor.PrimeDerivativeFwd
 end
 
 # Some Basic Mixed Mode tests
+# TODO: unbreak this
 function sin_twice_fwd(x)
     let var"'" = Diffractor.PrimeDerivativeFwd
             sin''(x)
@@ -90,4 +90,46 @@ end
     end
 end
 
+
+@testset "structs" begin
+    struct IDemo
+        x::Float64
+        y::Float64
+    end
+
+    function foo(a)
+        obj = IDemo(2.0, a)
+        return obj.x * obj.y
+    end
+
+    let var"'" = Diffractor.PrimeDerivativeFwd
+        @test foo'(100.0) == 2.0
+        @test foo''(100.0) == 0.0
+    end
+end
+
+@testset "tuples" begin
+    function foo(a)
+        tup = (2.0, a)
+        return first(tup) * tup[2]
+    end
+
+    let var"'" = Diffractor.PrimeDerivativeFwd
+        @test foo'(100.0) == 2.0
+        @test foo''(100.0) == 0.0
+    end
+end
+
+@testset "vararg" begin
+    function foo(a)
+        tup = (2.0, a)
+        return *(tup...)
+    end
+
+    let var"'" = Diffractor.PrimeDerivativeFwd
+        @test foo'(100.0) == 2.0
+        @test foo''(100.0) == 0.0
+    end
+end
+
 end