Merge pull request #153 from JuliaDiff/ox/tangentforward

Fix getindex on CompositeBundle of a struct

Author: oxinabox

src/tangent.jl

@@ -296,12 +296,9 @@ 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()
-    Tangent{B}(map(tb.tup) do el
-        el[tti]
-    end...)
+    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)...)

test/tangent.jl

@@ -1,6 +1,8 @@
 module tagent
 using Diffractor
 using Diffractor: AbstractZeroBundle, ZeroBundle, DNEBundle
+using Diffractor: TaylorBundle, TaylorTangentIndex, CompositeBundle
+using ChainRulesCore
 using Test
 
 @testset "AbstractZeroBundle" begin
@@ -24,4 +26,22 @@ using Test
     end
 end
 
+@testset "AD through constructor" begin
+    #https://github.com/JuliaDiff/Diffractor.jl/issues/152
+    # hits `getindex(::CompositeBundle{Foo152}, ::TaylorTangentIndex)`
+    struct Foo152
+        x::Float64
+    end
+
+    # Unit Test
+    cb = CompositeBundle{1, Foo152}((TaylorBundle{1, Float64}(23.5, (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)
+end
+
 end  # module