Add tests for some rrules

Author: tansongchen

benchmark/pinn.jl

@@ -15,7 +15,7 @@ model = Chain(
 trial(model, x) = x[1] * (1 - x[1]) * x[2] * (1 - x[2]) * model(x)
 
 M = 100
-data = [rand(input) for _ in 1:M]
+data = [rand(Float32, input) for _ in 1:M]
 function loss_by_finitediff(model, x)
     ε = cbrt(eps(Float32))
     ε₁ = [ε, 0]
@@ -27,7 +27,7 @@ function loss_by_finitediff(model, x)
 end
 function loss_by_taylordiff(model, x)
     f(x) = trial(model, x)
-    error = derivative(f, x, [1., 0.], 2) + derivative(f, x, [0., 1.], 2) + sin(π * x[1]) * sin(π * x[2])
+    error = derivative(f, x, Float32[1, 0], 2) + derivative(f, x, Float32[0, 1], 2) + sin(π * x[1]) * sin(π * x[2])
     abs2(error)
 end
 

src/chainrules.jl

@@ -38,34 +38,35 @@ function rrule(::Type{TaylorScalar{T, N}}, v::NTuple{N, T}) where {N, T <: Numbe
     return TaylorScalar(v), taylor_scalar_pullback
 end
 
-function rrule(::typeof(value), t::TaylorScalar)
-    value_pullback(v̄::NTuple) = NoTangent(), TaylorScalar(v̄)
+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(Tuple(v̄))
+    value_pullback(v̄) = NoTangent(), TaylorScalar(map(x -> convert(T, x), Tuple(v̄)))
     return value(t), value_pullback
 end
 
 function rrule(::typeof(extract_derivative), t::TaylorScalar{T, N},
                i::Integer) where {N, T <: Number}
     function extract_derivative_pullback(d̄)
-        NoTangent(), TaylorScalar((zeros(T, i - 1)..., d̄, zeros(T, N - i)...)), NoTangent()
+        NoTangent(), TaylorScalar{T, N}(ntuple(j -> j === i ? d̄ : zero(T), Val(N))), NoTangent()
     end
     return extract_derivative(t, i), extract_derivative_pullback
 end
 
 function rrule(::typeof(*), A::Matrix{S}, t::Vector{TaylorScalar{T, N}}) where {N, S <: Number, T}
     project_A = ProjectTo(A)
-    gemv_pullback(x̄) = NoTangent(), project_A(contract.(x̄, transpose(t))), transpose(A) * x̄
+    gemv_pullback(x̄) = NoTangent(), @thunk(project_A(contract.(x̄, transpose(t)))), @thunk(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̄
+    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̄, map(primal, x̄)
+    vadd_pullback(x̄) = NoTangent(), x̄, ProjectTo(v)(x̄)
     return t + v, vadd_pullback
 end
 

src/primitive.jl

@@ -13,7 +13,7 @@ import Base: hypot, max, min
 
 @inline sqrt(t::TaylorScalar) = t^0.5
 @inline cbrt(t::TaylorScalar) = ^(t, 1 / 3)
-@inline inv(t::TaylorScalar) = 1 / t
+@inline inv(t::TaylorScalar) = one(t) / t
 @inline abs(t::TaylorScalar) = primal(t) >= 0 ? t : -t
 
 for func in (:exp, :expm1, :exp2, :exp10)

test/zygote.jl

@@ -1,8 +1,24 @@
 using Zygote
 
-@testset "Zygote compatibility" begin @test gradient(x -> derivative(x -> x * x, x, 1),
+@testset "Zygote for mixed derivative" begin
+    some_number = 0.7
+    for f in (exp, log, sqrt, sin, asin, sinh, asinh)
+        @test gradient(x -> derivative(f, x, 2), some_number)[1] ≈ derivative(f, some_number, 3)
+    end
+    @test gradient(x -> derivative(x -> x * x, x, 1),
                                                      5.0)[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
+
+@testset "Zygote for parameter optimization" begin
+    linear_model(x, p) = exp.(p * x)[1]
+    some_x, some_v, some_p = [.58, .36], [.23, .11], [.49 .96]
+    loss_taylor(p) = derivative(x -> linear_model(x, p), some_x, some_v, 1)
+    ε = cbrt(eps(Float64))
+    loss_finite(p) = let f = x -> linear_model(x, p)
+        (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]
+end