Improve coverage

Author: tansongchen

.github/workflows/CI.yml

@@ -3,7 +3,7 @@ on:
   push:
     branches:
       - main
-    tags: '*'
+    tags: ['*']
   pull_request:
 concurrency:
   # Skip intermediate builds: always.
@@ -36,6 +36,7 @@ jobs:
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v2
         with:
+          token: ${{ secrets.CODECOV_TOKEN }}
           files: lcov.info
   docs:
     name: Documentation

benchmark/linearmodel.jl

@@ -12,7 +12,7 @@ end
 data = rand(2)
 
 function loss(model)
-    derivative(model, data, [1., 0.], Val(2))
+    derivative(model, data, [1.0, 0.0], Val(2))
 end
 
 model = LinearModel(rand(1, 2))

benchmark/pinn.jl

@@ -6,12 +6,10 @@ using Plots
 const input = 2
 const hidden = 16
 
-model = Chain(
-    Dense(input => hidden, sin),
-    Dense(hidden => hidden, sin),
-    Dense(hidden => 1),
-    first
-)
+model = Chain(Dense(input => hidden, sin),
+              Dense(hidden => hidden, sin),
+              Dense(hidden => 1),
+              first)
 trial(model, x) = x[1] * (1 - x[1]) * x[2] * (1 - x[2]) * model(x)
 
 M = 100
@@ -27,7 +25,8 @@ function loss_by_finitediff(model, x)
 end
 function loss_by_taylordiff(model, x)
     f(x) = trial(model, x)
-    error = derivative(f, x, Float32[1, 0], 2) + derivative(f, x, Float32[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

@@ -1,25 +1,23 @@
 import ChainRulesCore: rrule, RuleConfig, ProjectTo
 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,
-    inv, sqrt, cbrt,
-    sin, cos, tan, cot, sec, csc,
-    asin, acos, atan, acot, asec, acsc,
-    sinh, cosh, tanh, coth, sech, csch,
-    asinh, acosh, atanh, acoth, asech, acsch,
-]
+function contract(a::TaylorScalar{T, N}, b::TaylorScalar{S, N}) where {T, S, N}
+    mapreduce(*, +, value(a), value(b))
+end
+
+NONLINEAR_UNARY_FUNCTIONS = Function[exp, exp2, exp10, expm1,
+                                     log, log2, log10, log1p,
+                                     inv, sqrt, cbrt,
+                                     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[
-    *, /, ^
-]
+NONLINEAR_BINARY_FUNCTIONS = Function[*, /, ^]
 
 for func in NONLINEAR_BINARY_FUNCTIONS
     @eval @opt_out rrule(::typeof($func), ::TaylorScalar, ::TaylorScalar)
@@ -49,31 +47,43 @@ end
 function rrule(::typeof(extract_derivative), t::TaylorScalar{T, N},
                i::Integer) where {N, T <: Number}
     function extract_derivative_pullback(d̄)
-        NoTangent(), TaylorScalar{T, N}(ntuple(j -> j === i ? d̄ : zero(T), Val(N))), 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}
+function rrule(::typeof(*), A::Matrix{S},
+               t::Vector{TaylorScalar{T, N}}) where {N, S <: Number, T}
     project_A = ProjectTo(A)
-    gemv_pullback(x̄) = NoTangent(), @thunk(project_A(contract.(x̄, transpose(t)))), @thunk(transpose(A) * x̄)
+    function gemv_pullback(x̄)
+        NoTangent(), @thunk(project_A(contract.(x̄, transpose(t)))), @thunk(transpose(A)*x̄)
+    end
     return A * t, gemv_pullback
 end
 
-function rrule(::typeof(+), v::Vector{T}, t::Vector{TaylorScalar{T, N}}) where {N, T <: Number}
+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}
+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 +(t::Vector{TaylorScalar{T, N}}, v::Vector{T}) where {N, T <: Number} = t + v, x̄ -> (x̄, map(primal, x̄))
+@adjoint function +(t::Vector{TaylorScalar{T, N}}, v::Vector{T}) where {N, T <: Number}
+    t + v, x̄ -> (x̄, map(primal, x̄))
+end
 
-@adjoint +(v::Vector{T}, t::Vector{TaylorScalar{T, N}}) where {N, T <: Number} = v + t, x̄ -> (map(primal, x̄), x̄)
+@adjoint function +(v::Vector{T}, t::Vector{TaylorScalar{T, N}}) where {N, T <: Number}
+    v + t, x̄ -> (map(primal, x̄), x̄)
+end
 
 (project::ProjectTo{T})(dx::TaylorScalar{T, N}) where {N, T <: Number} = primal(dx)
 
-(project::ProjectTo{S})(dx::TaylorScalar{T, N}) where {N, T <: Number, S <: Real} = project(primal(dx))
+function (project::ProjectTo{S})(dx::TaylorScalar{T, N}) where {N, T <: Number, S <: Real}
+    project(primal(dx))
+end

test/zygote.jl

@@ -3,22 +3,24 @@ using Zygote
 @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)
+        @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
+                   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.]
+    @test gradient(x -> derivative(g, x, [1.0, 0.0], 1), [1.0, 2.0])[1] ≈ [2.0, 0.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]
+    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)
     ε = 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
+    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