Fix gradient for gemv and vadd

Author: tansongchen

Project.toml

@@ -8,10 +8,12 @@ ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChainRulesOverloadGeneration = "f51149dc-2911-5acf-81fc-2076a2a81d4f"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
+ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [compat]
 ChainRules = "1"
 ChainRulesCore = "1"
 ChainRulesOverloadGeneration = "0.1"
 SymbolicUtils = "0.19, 0.20"
+ZygoteRules = "0.2"
 julia = "1.6"

benchmark/linearmodel.jl

@@ -0,0 +1,19 @@
+using TaylorDiff, Zygote
+using LinearAlgebra: dot
+
+sigmoid(x) = (one(x) + tanh(x)) / 2
+
+struct LinearModel
+    c::Matrix
+end
+
+(m::LinearModel)(x) = sigmoid((m.c * x)[1])
+
+data = rand(2)
+
+function loss(model)
+    derivative(model, data, [1., 0.], Val(2))
+end
+
+model = LinearModel(rand(1, 2))
+gradient(loss, model)

benchmark/pinn.jl

@@ -1,4 +1,4 @@
-using TaylorDiff
+using TaylorDiff, Zygote
 
 const input = 2
 const hidden = 16
@@ -10,7 +10,7 @@ struct PINN
     b₂
 end
 
-(pinn::PINN)(x) = x[1] * (1 - x[1]) * x[2] * (1 - x[2]) * first(pinn.W₂ * (pinn.W₁ * x + pinn.b₁) + pinn.b₂)
+(pinn::PINN)(x) = x[1] * (1 - x[1]) * x[2] * (1 - x[2]) * first(pinn.W₂ * exp.(pinn.W₁ * x + pinn.b₁) + pinn.b₂)
 
 dataset = [rand(input) for i in 1:10]
 function loss(pinn)
@@ -21,12 +21,6 @@ function loss(pinn)
     out
 end
 
-# function simple(w)
-#     derivative(x -> sum(w * x), [0.5, 0.7], [1., 0.], Val(2))
-# end
-# w = rand(hidden, input)
-# gradient(simple, w)
-
 myPINN = PINN(rand(hidden, input), rand(hidden), rand(1, hidden), rand(1))
 
 gradient(loss, myPINN)

src/chainrules.jl

@@ -1,9 +1,33 @@
 import ChainRulesCore: rrule, RuleConfig
+using ZygoteRules: @adjoint
+
+contract(a::TaylorScalar{T, N}, b::TaylorScalar{S, N}) where {T, S, N} = mapreduce(*, +, value(a), value(b))
+
+NONLINEAR_UNARY_FUNCTIONS = Function[
+    exp, exp2, exp10, expm1,
+    log, log2, log10, log1p,
+    sin, cos, tan, cot, sec, csc,
+    asin, acos, atan, acot, asec, acsc,
+    sinh, cosh, tanh, coth, sech, csch,
+    asinh, acosh, atanh, acoth, asech, acsch,
+]
+
+for func in NONLINEAR_UNARY_FUNCTIONS
+    @eval @opt_out rrule(::typeof($func), ::TaylorScalar)
+end
+
+NONLINEAR_BINARY_FUNCTIONS = Function[
+    *, /, ^
+]
+
+for func in NONLINEAR_BINARY_FUNCTIONS
+    @eval @opt_out rrule(::typeof($func), ::TaylorScalar, ::TaylorScalar)
+    @eval @opt_out rrule(::typeof($func), ::TaylorScalar, ::Number)
+    @eval @opt_out rrule(::typeof($func), ::Number, ::TaylorScalar)
+end
+
+# Other special cases
 
-@opt_out rrule(::Any, ::TaylorScalar, ::TaylorScalar)
-@opt_out rrule(::Any, ::TaylorScalar, ::Any)
-@opt_out rrule(::typeof(*), ::TaylorScalar, ::TaylorScalar)
-@opt_out rrule(::typeof(^), ::TaylorScalar, ::Any)
 @opt_out rrule(::typeof(Base.literal_pow), ::typeof(^), x::TaylorScalar, ::Val{p}) where {p}
 @opt_out rrule(::RuleConfig, ::typeof(Base.literal_pow), ::typeof(^), x::TaylorScalar,
                ::Val{p}) where {p}
@@ -29,7 +53,7 @@ function rrule(::typeof(extract_derivative), t::TaylorScalar{T, N},
 end
 
 function rrule(::typeof(*), A::Matrix{T}, t::Vector{TaylorScalar{T, N}}) where {N, T <: Number}
-    gemv_pullback(x̄) = NoTangent(), map(primal, x̄) * transpose(map(primal, t)), transpose(A) * x̄
+    gemv_pullback(x̄) = NoTangent(), contract.(x̄, transpose(t)), transpose(A) * x̄
     return A * t, gemv_pullback
 end
 
@@ -42,3 +66,7 @@ function rrule(::typeof(+), t::Vector{TaylorScalar{T, N}}, v::Vector{T}) where {
     vadd_pullback(x̄) = NoTangent(), x̄, map(primal, x̄)
     return t + v, vadd_pullback
 end
+
+@adjoint +(t::Vector{TaylorScalar{T, N}}, v::Vector{T}) where {N, T <: Number} = t + v, x̄ -> (x̄, map(primal, x̄))
+
+@adjoint +(v::Vector{T}, t::Vector{TaylorScalar{T, N}}) where {N, T <: Number} = v + t, x̄ -> (map(primal, x̄), x̄)