Improve coverage

Author: tansongchen

src/chainrules.jl

@@ -38,7 +38,6 @@ end
 
 function rrule(::typeof(value), t::TaylorScalar{T, N}) where {N, T}
     value_pullback(v̄::NTuple{N, T}) = NoTangent(), TaylorScalar(v̄)
-    value_pullback(v̄::Tuple) = NoTangent(), TaylorScalar(map(x -> convert(T, x), v̄))
     # for structural tangent, convert to tuple
     value_pullback(v̄) = NoTangent(), TaylorScalar(map(x -> convert(T, x), Tuple(v̄)))
     return value(t), value_pullback
@@ -62,28 +61,14 @@ function rrule(::typeof(*), A::Matrix{S},
     return A * t, gemv_pullback
 end
 
-function rrule(::typeof(+), v::Vector{T},
-               t::Vector{TaylorScalar{T, N}}) where {N, T <: Number}
-    vadd_pullback(x̄) = NoTangent(), ProjectTo(v)(x̄), x̄
-    return v + t, vadd_pullback
-end
-
-function rrule(::typeof(+), t::Vector{TaylorScalar{T, N}},
-               v::Vector{T}) where {N, T <: Number}
-    vadd_pullback(x̄) = NoTangent(), x̄, ProjectTo(v)(x̄)
-    return t + v, vadd_pullback
-end
-
 @adjoint function +(t::Vector{TaylorScalar{T, N}}, v::Vector{T}) where {N, T <: Number}
-    t + v, x̄ -> (x̄, map(primal, x̄))
+    project_v = ProjectTo(v)
+    t + v, x̄ -> (x̄, project_v(x̄))
 end
 
 @adjoint function +(v::Vector{T}, t::Vector{TaylorScalar{T, N}}) where {N, T <: Number}
-    v + t, x̄ -> (map(primal, x̄), x̄)
+    project_v = ProjectTo(v)
+    v + t, x̄ -> (project_v(x̄), x̄)
 end
 
 (project::ProjectTo{T})(dx::TaylorScalar{T, N}) where {N, T <: Number} = primal(dx)
-
-function (project::ProjectTo{S})(dx::TaylorScalar{T, N}) where {N, T <: Number, S <: Real}
-    project(primal(dx))
-end

src/primitive.jl

@@ -14,7 +14,6 @@ import Base: hypot, max, min
 @inline sqrt(t::TaylorScalar) = t^0.5
 @inline cbrt(t::TaylorScalar) = ^(t, 1 / 3)
 @inline inv(t::TaylorScalar) = one(t) / t
-@inline abs(t::TaylorScalar) = primal(t) >= 0 ? t : -t
 
 for func in (:exp, :expm1, :exp2, :exp10)
     @eval @generated function $func(t::TaylorScalar{T, N}) where {T, N}

test/zygote.jl

@@ -14,12 +14,14 @@ using Zygote
 end
 
 @testset "Zygote for parameter optimization" begin
-    linear_model(x, p) = exp.(p * x)[1]
-    some_x, some_v, some_p = [0.58, 0.36], [0.23, 0.11], [0.49 0.96]
-    loss_taylor(p) = derivative(x -> linear_model(x, p), some_x, some_v, 1)
+    gradient(p -> derivative(x -> sum(exp.(x + p)), [1.0, 1.0], [1.0, 0.0], 1), [0.5, 0.7])
+    gradient(p -> derivative(x -> sum(exp.(p + x)), [1.0, 1.0], [1.0, 0.0], 1), [0.5, 0.7])
+    linear_model(x, p, b) = exp.(b + p * x + b)[1]
+    some_x, some_v, some_p, some_b = [0.58, 0.36], [0.23, 0.11], [0.49 0.96], [0.88]
+    loss_taylor(p) = derivative(x -> linear_model(x, p, some_b), some_x, some_v, 1)
     ε = cbrt(eps(Float64))
     loss_finite(p) =
-        let f = x -> linear_model(x, p)
+        let f = x -> linear_model(x, p, some_b)
             (f(some_x + ε * some_v) - f(some_x - ε * some_v)) / 2ε
         end
     @test gradient(loss_taylor, some_p)[1] ≈ gradient(loss_finite, some_p)[1]