Add rrule for gemv

Author: tansongchen

benchmark/pinn.jl

@@ -0,0 +1,32 @@
+using TaylorDiff
+
+const input = 2
+const hidden = 16
+
+struct PINN
+    W₁
+    b₁
+    W₂
+    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₂)
+
+dataset = [rand(input) for i in 1:10]
+function loss(pinn)
+    out = 0.0
+    for x in dataset
+        out += derivative(pinn, x, [1., 0.], Val(2))
+    end
+    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,6 +1,5 @@
 import ChainRulesCore: rrule, RuleConfig
 
-@opt_out rrule(::Any, ::TaylorScalar)
 @opt_out rrule(::Any, ::TaylorScalar, ::TaylorScalar)
 @opt_out rrule(::Any, ::TaylorScalar, ::Any)
 @opt_out rrule(::typeof(*), ::TaylorScalar, ::TaylorScalar)
@@ -28,3 +27,18 @@ function rrule(::typeof(extract_derivative), t::TaylorScalar{T, N},
     end
     return extract_derivative(t, i), extract_derivative_pullback
 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̄
+    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(), map(primal, 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̄, map(primal, x̄)
+    return t + v, vadd_pullback
+end

src/scalar.jl

@@ -55,6 +55,7 @@ end
 @inline value(t::TaylorScalar) = t.value
 @inline extract_derivative(t::TaylorScalar, i::Integer) = t.value[i]
 @inline extract_derivative(r, i::Integer) = false
+@inline primal(t::TaylorScalar) = extract_derivative(t, 1)
 
 @inline zero(::Type{TaylorScalar{T, N}}) where {T, N} = TaylorScalar{T, N}(zero(T))
 @inline one(::Type{TaylorScalar{T, N}}) where {T, N} = TaylorScalar{T, N}(one(T))

test/zygote.jl

@@ -2,5 +2,7 @@ using Zygote
 
 @testset "Zygote compatibility" begin @test gradient(x -> derivative(x -> x * x, x, 1),
                                                      5.0)[1] ≈ 2.0
-    # @test gradient(x -> derivative(g, x, 1, 1), [1., 2.])[1] ≈ [2., 0.]
+
+    g(x) = x[1] * x[1] + x[2] * x[2]
+    @test gradient(x -> derivative(g, x, [1., 0.], 1), [1., 2.])[1] ≈ [2., 0.]
 end