fix and test taylor_compatible

Author: oxinabox

src/stage1/forward.jl

@@ -55,16 +55,16 @@ function taylor_compatible(a::ATB{N}, b::ATB{N}) where {N}
 end
 
 function taylor_compatible(r::TaylorBundle{N, Tuple{B1,B2}}) where {N, B1,B2}
-    partial(r, 1)[1] = primal(r)[2] || return false
-    return all(1:N-1) do ii
+    partial(r, 1)[1] == primal(r)[2] || return false
+    return all(1:N-1) do i
         partial(r, i+1)[1] == partial(r, i)[2]
     end
 end
 function shuffle_up(r::TaylorBundle{N, Tuple{B1,B2}}) where {N, B1,B2}
     the_primal = primal(r)[1]
     if taylor_compatible(r)
         the_partials = ntuple(N+1) do i
-            if ii <= N
+            if i <= N
                 partial(r, i)[1]  # == `partial(r,i-1)[2]` (except first which is primal(r)[2])
             else  # ii = N+1
                 partial(r, i-1)[2]

test/forward.jl

@@ -1,16 +1,12 @@
 module forward_tests
 using Diffractor
-using Diffractor: var"'", ∂⃖, DiffractorRuleConfig, ZeroBundle
+using Diffractor: TaylorBundle
 using ChainRules
 using ChainRulesCore
 using ChainRulesCore: ZeroTangent, NoTangent, frule_via_ad, rrule_via_ad
 using LinearAlgebra
-
 using Test
 
-const fwd = Diffractor.PrimeDerivativeFwd
-const bwd = Diffractor.PrimeDerivativeBack
-
 
 
 # Minimal 2-nd order forward smoke test
@@ -131,4 +127,41 @@ end
     end
 end
 
+
+@testset "taylor_compatible" begin
+    taylor_compatible = Diffractor.taylor_compatible
+
+    @test taylor_compatible(
+        TaylorBundle{1}(10.0, (20.0,)),
+        TaylorBundle{1}(20.0, (30.0,))
+    )
+    @test !taylor_compatible(
+        TaylorBundle{1}(10.0, (20.0,)),
+        TaylorBundle{1}(21.0, (30.0,))
+    )
+    @test taylor_compatible(
+        TaylorBundle{2}(10.0, (20.0, 30.)),
+        TaylorBundle{2}(20.0, (30.0, 40.))
+    )
+    @test !taylor_compatible(
+        TaylorBundle{2}(10.0, (20.0, 30.0)),
+        TaylorBundle{2}(20.0, (31.0, 40.0))
+    )
+
+
+    tuptan(args...) = Tangent{typeof(args)}(args...)
+    @test taylor_compatible(
+        TaylorBundle{1}((10.0, 20.0), (tuptan(20.0, 30.0),)),
+    )
+    @test taylor_compatible(
+        TaylorBundle{2}((10.0, 20.0), (tuptan(20.0, 30.0),tuptan(30.0, 40.0))),
+    )
+    @test !taylor_compatible(
+        TaylorBundle{1}((10.0, 20.0), (tuptan(21.0, 30.0),)),
+    )
+    @test !taylor_compatible(
+        TaylorBundle{2}((10.0, 20.0), (tuptan(20.0, 31.0),tuptan(30.0, 40.0))),
+    )
+end
+
 end