Remove most uses of CompositeTangent

Author: oxinabox

src/AbstractDifferentiation.jl

@@ -10,19 +10,7 @@ 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::AbstractArray{<:Number}) = TaylorBundle{1}(x, (dx,))
-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
-
+bundle(x, dx) = TaylorBundle{1}(x, (dx,))
 
 AD.@primitive function pushforward_function(b::DiffractorForwardBackend, f, args...)
     return function pushforward(vs)

src/stage1/forward.jl

@@ -4,14 +4,6 @@ partial(x::TaylorTangent, i) = getfield(getfield(x, :coeffs), i)
 partial(x::UniformTangent, i) = getfield(x, :val)
 partial(x::ProductTangent, i) = ProductTangent(map(x->partial(x, i), getfield(x, :factors)))
 partial(x::AbstractZero, i) = x
-partial(x::CompositeBundle{N, B}, i) where {N, B<:Tuple} = Tangent{B}(map(x->partial(x, i), getfield(x, :tup))...)
-function partial(x::CompositeBundle{N, B}, i) where {N, B}
-    # This is tangent for a struct, but fields partials are each stored in a plain tuple
-    # so we add the names back using the primal `B`
-    # TODO: If required this can be done as a `@generated` function so it is type-stable
-    backing = NamedTuple{fieldnames(B)}(map(x->partial(x, i), getfield(x, :tup)))
-    return Tangent{B, typeof(backing)}(backing)
-end
 
 
 primal(x::AbstractTangentBundle) = x.primal
@@ -42,14 +34,6 @@ function shuffle_down(b::TaylorBundle{N, B}) where {N, B}
         ntuple(_sdown, N-1))
 end
 
-function shuffle_down(b::CompositeBundle{N, B}) where {N, B}
-    z = CompositeBundle{N-1, CompositeBundle{1, B}}(
-        (CompositeBundle{N-1, Tuple}(
-            map(shuffle_down, b.tup)
-        ),)
-    )
-    z
-end
 
 function shuffle_up(r::TaylorBundle{1, Tuple{B1,B2}}) where {B1,B2}
     z₀ = primal(r)[1]
@@ -63,18 +47,7 @@ function shuffle_up(r::TaylorBundle{1, Tuple{B1,B2}}) where {B1,B2}
     end
 end
 
-function shuffle_up(r::CompositeBundle{1})
-    z₀ = primal(r.tup[1])
-    z₁ = partial(r.tup[1], 1)
-    z₂ = primal(r.tup[2])
-    z₁₂ = partial(r.tup[2], 1)
-    if z₁ == z₂
-        return TaylorBundle{2}(z₀, (z₁, z₁₂))
-    else
-        return ExplicitTangentBundle{2}(z₀, (z₁, z₂, z₁₂))
-    end
-end
-
+#==
 function taylor_compatible(a::ATB{N}, b::ATB{N}) where {N}
     primal(b) === a[TaylorTangentIndex(1)] || return false
     return all(1:(N-1)) do i
@@ -88,7 +61,7 @@ isswifty(::UniformBundle) = true
 isswifty(b::CompositeBundle) = all(isswifty, b.tup)
 isswifty(::Any) = false
 
-#TODO: port this to TaylorTangent:
+#TODO: port this to TaylorTangent over composite structures
 function shuffle_up(r::CompositeBundle{N}) where {N}
     a, b = r.tup
     if isswifty(a) && isswifty(b) && taylor_compatible(a, b)
@@ -102,6 +75,7 @@ function shuffle_up(r::CompositeBundle{N}) where {N}
             ntuple(i->partial(b,i), 1<<(N+1)-1)...))
     end
 end
+==#
 
 function shuffle_up(r::UniformBundle{N, B, U}) where {N, B, U}
     (a, b) = primal(r)
@@ -198,13 +172,6 @@ end
         map(y->lifted_getfield(y, s), x.tangent.coeffs))
 end
 
-@Base.constprop :aggressive function (::∂☆{N})(::ATB{N, typeof(getfield)}, x::CompositeBundle{N}, s::AbstractTangentBundle{N, Int}) where {N}
-    x.tup[primal(s)]
-end
-
-@Base.constprop :aggressive function (::∂☆{N})(::ATB{N, typeof(getfield)}, x::CompositeBundle{N, B}, s::AbstractTangentBundle{N, Symbol}) where {N, B}
-    x.tup[Base.fieldindex(B, primal(s))]
-end
 
 @Base.constprop :aggressive function (::∂☆{N})(f::ATB{N, typeof(getfield)}, x::UniformBundle{N, <:Any, U}, s::AbstractTangentBundle{N}) where {N, U}
     UniformBundle{N,<:Any,U}(getfield(primal(x), primal(s)), x.tangent.val)
@@ -223,9 +190,12 @@ struct FwdMap{N, T<:AbstractTangentBundle{N}}
 end
 (f::FwdMap{N})(args::AbstractTangentBundle{N}...) where {N} = ∂☆{N}()(f.f, args...)
 
+#==
+# TODO port this to TaylorBundle over composite structure
 function (::∂☆{N})(::ZeroBundle{N, typeof(map)}, f::ATB{N}, tup::CompositeBundle{N, <:Tuple}) where {N}
     ∂vararg{N}()(map(FwdMap(f), tup.tup)...)
 end
+==#
 
 function (::∂☆{N})(::ZeroBundle{N, typeof(map)}, f::ATB{N}, args::ATB{N, <:AbstractArray}...) where {N}
     # TODO: This could do an inplace map! to avoid the extra rebundling
@@ -267,23 +237,28 @@ function (this::∂☆{N})(::ZeroBundle{N, typeof(Core._apply_iterate)}, iterate
     Core._apply_iterate(FwdIterate(iterate), this, (f,), args...)
 end
 
+#==
+#TODO: port this to TaylorTangent over composite structures
 function (this::∂☆{N})(::ZeroBundle{N, typeof(iterate)}, t::CompositeBundle{N, <:Tuple}) where {N}
     r = iterate(t.tup)
     r === nothing && return ZeroBundle{N}(nothing)
     ∂vararg{N}()(r[1], ZeroBundle{N}(r[2]))
 end
 
+#TODO: port this to TaylorTangent over composite structures
 function (this::∂☆{N})(::ZeroBundle{N, typeof(iterate)}, t::CompositeBundle{N, <:Tuple}, a::ATB{N}, args::ATB{N}...) where {N}
     r = iterate(t.tup, primal(a), map(primal, args)...)
     r === nothing && return ZeroBundle{N}(nothing)
     ∂vararg{N}()(r[1], ZeroBundle{N}(r[2]))
 end
 
+#TODO: port this to TaylorTangent over composite structures
 function (this::∂☆{N})(::ZeroBundle{N, typeof(Base.indexed_iterate)}, t::CompositeBundle{N, <:Tuple}, i::ATB{N}) where {N}
     r = Base.indexed_iterate(t.tup, primal(i))
     ∂vararg{N}()(r[1], ZeroBundle{N}(r[2]))
 end
 
+#TODO: port this to TaylorTangent over composite structures
 function (this::∂☆{N})(::ZeroBundle{N, typeof(Base.indexed_iterate)}, t::CompositeBundle{N, <:Tuple}, i::ATB{N}, st1::ATB{N}, st::ATB{N}...) where {N}
     r = Base.indexed_iterate(t.tup, primal(i), primal(st1), map(primal, st)...)
     ∂vararg{N}()(r[1], ZeroBundle{N}(r[2]))
@@ -293,10 +268,11 @@ function (this::∂☆{N})(::ZeroBundle{N, typeof(Base.indexed_iterate)}, t::Tan
     ∂vararg{N}()(this(ZeroBundle{N}(getfield), t, i), ZeroBundle{N}(primal(i) + 1))
 end
 
-
+#TODO: port this to TaylorTangent over composite structures
 function (this::∂☆{N})(::ZeroBundle{N, typeof(getindex)}, t::CompositeBundle{N, <:Tuple}, i::ZeroBundle) where {N}
     t.tup[primal(i)]
 end
+==#
 
 function (this::∂☆{N})(::ZeroBundle{N, typeof(typeof)}, x::ATB{N}) where {N}
     DNEBundle{N}(typeof(primal(x)))

src/tangent.jl

@@ -290,33 +290,6 @@ end
 
 Base.getindex(u::UniformBundle, ::TaylorTangentIndex) = u.tangent.val
 
-"""
-    CompositeBundle{N, B, B <: Tuple}
-
-Represents the tagent bundle where the base space is some tuple or struct type.
-Mathematically, this tangent bundle is the product bundle of the individual
-element bundles.
-"""
-struct CompositeBundle{N, B, T<:Tuple{Vararg{AbstractTangentBundle{N}}}} <: AbstractTangentBundle{N, B}
-    tup::T
-end
-CompositeBundle{N, B}(tup::T) where {N, B, T} = CompositeBundle{N, B, T}(tup)
-
-function Base.getindex(tb::CompositeBundle{N, B} where N, tti::TaylorTangentIndex) where {B}
-    B <: SArray && error()
-    return partial(tb, tti.i)
-end
-
-primal(b::CompositeBundle{N, <:Tuple} where N) = map(primal, b.tup)
-function primal(b::CompositeBundle{N, T} where N) where T<:CompositeBundle
-    T(map(primal, b.tup)...)
-end
-@generated primal(b::CompositeBundle{N, B} where N) where {B} =
-    quote
-        x = map(primal, b.tup)
-        $(Expr(:splatnew, B, :x))
-    end
-
 expand_singleton_to_array(asize, a::AbstractZero) = fill(a, asize...)
 expand_singleton_to_array(asize, a::AbstractArray) = a
 

test/AbstractDifferentiationTests.jl

@@ -7,14 +7,14 @@ backend = Diffractor.DiffractorForwardBackend()
 
     @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.5=>2.5, Tangent{Pair{Float64, Float64}}(first=1.0, second=2.0)) isa Diffractor.TaylorBundle{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}
+    @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.TaylorBundle{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 b isa Diffractor.TaylorBundle{1}
     @test iszero(t.first)
     @test t.second.first == 1.0
     @test t.second.second == 2.0

test/tangent.jl

@@ -1,7 +1,7 @@
 module tagent
 using Diffractor
 using Diffractor: AbstractZeroBundle, ZeroBundle, DNEBundle
-using Diffractor: TaylorBundle, TaylorTangentIndex, CompositeBundle
+using Diffractor: TaylorBundle, TaylorTangentIndex
 using Diffractor: ExplicitTangent, TaylorTangent, truncate
 using ChainRulesCore
 using Test
@@ -28,21 +28,22 @@ using Test
 end
 
 @testset "AD through constructor" begin
-    #https://github.com/JuliaDiff/Diffractor.jl/issues/152
-    # hits `getindex(::CompositeBundle{Foo152}, ::TaylorTangentIndex)`
+    # https://github.com/JuliaDiff/Diffractor.jl/issues/152
+    # Though we have now removed the underlying cause, we keep this as a regression test just in case
     struct Foo152
         x::Float64
     end
 
     # Unit Test
-    cb = CompositeBundle{1, Foo152}((TaylorBundle{1, Float64}(23.5, (1.0,)),))
+    cb = TaylorBundle{1, Foo152}(Foo152(23.5), (Tangent{Foo152}(;x=1.0),))   
     tti = TaylorTangentIndex(1,)
     @test cb[tti] == Tangent{Foo152}(; x=1.0)
 
     # Integration  Test
-    var"'" = Diffractor.PrimeDerivativeFwd
-    f(x) = Foo152(x)
-    @test f'(23.5) == Tangent{Foo152}(; x=1.0)
+    let var"'" = Diffractor.PrimeDerivativeFwd
+        f(x) = Foo152(x)
+        @test f'(23.5) == Tangent{Foo152}(; x=1.0)
+    end
 end
 
 @testset "truncate" begin