Address code review comments

Author: YingboMa

src/rulesets/LinearAlgebra/blas.jl

@@ -37,7 +37,7 @@ end
 
 function frule(::typeof(BLAS.nrm2), x, _, Δ)
     Ω = BLAS.nrm2(x)
-    return Ω, sum(Δx * cast(@thunk(x * inv(Ω))))
+    return Ω, sum(Δx .* @thunk(x * inv(Ω)))
 end
 
 function rrule(::typeof(BLAS.nrm2), x)
@@ -68,12 +68,15 @@ end
 #####
 
 function frule(::typeof(BLAS.asum), x, _, Δx)
-    return BLAS.asum(x), sum(cast(sign, x) * Δx)
+    return BLAS.asum(x), sum(zip(x, Δx)) do xs
+        x, Δx = xs
+        return sign(x) * Δx
+    end
 end
 
 function rrule(::typeof(BLAS.asum), x)
     function asum_pullback(ΔΩ)
-        return (NO_FIELDS, @thunk(ΔΩ * cast(sign, x)))
+        return (NO_FIELDS, @thunk(ΔΩ * sign.(x)))
     end
     return BLAS.asum(x), asum_pullback
 end

src/rulesets/LinearAlgebra/dense.jl

@@ -25,15 +25,13 @@ end
 
 function frule(::typeof(inv), x::AbstractArray, _, Δx)
     Ω = inv(x)
-    m = @thunk(-Ω)
-    return Ω, m * Δx * Ω
+    return Ω, -Ω * Δx * Ω
 end
 
 function rrule(::typeof(inv), x::AbstractArray)
     Ω = inv(x)
-    m = @thunk(-Ω')
     function inv_pullback(ΔΩ)
-        return NO_FIELDS, m * ΔΩ * Ω'
+        return NO_FIELDS, -Ω' * ΔΩ * Ω'
     end
     return Ω, inv_pullback
 end

test/rulesets/Base/base.jl

@@ -52,24 +52,6 @@
             test_scalar(acotd, 1/x)
         end
         @testset "Multivariate" begin
-            @testset "atan2" begin
-                # https://en.wikipedia.org/wiki/Atan2
-                x, y = rand(2)
-                ratan = atan(x, y)
-                u = x^2 + y^2
-                datan = y/u - 2x/u
-
-                r, ṙ = frule(atan, x, y, Zero(), 1, 2)
-                @test r === ratan
-                @test ṙ === datan
-
-                r, pullback = rrule(atan, x, y)
-                @test r === ratan
-                dself, df1, df2 = pullback(1)
-                @test dself == NO_FIELDS
-                @test df1 + 2df2 === datan
-            end
-
             @testset "sincos" begin
                 x, Δx, x̄ = randn(3)
                 Δz = (randn(), randn())
@@ -91,11 +73,30 @@
             test_scalar(exp2, x)
             test_scalar(exp10, x)
 
+            test_scalar(cbrt, x)
+
+            if x >= 0
+                test_scalar(sqrt, x)
+                test_scalar(log, x)
+                test_scalar(log2, x)
+                test_scalar(log10, x)
+                test_scalar(log1p, x)
+            end
+        end
+    end
+
+    @testset "Unary complex functions" begin
+        for x in (-4.1, 6.4)
+            test_scalar(real, x)
+            test_scalar(imag, x)
+
+            test_scalar(abs, x)
+            test_scalar(hypot, x)
+
+            test_scalar(angle, x)
+            test_scalar(abs2, x)
             test_scalar(conj, x)
             test_scalar(adjoint, x)
-            test_scalar(abs2, x)
-
-            x isa Real && test_scalar(cbrt, x)
         end
     end
 

test/test_util.jl

@@ -26,15 +26,17 @@ function test_scalar(f, x; rtol=1e-9, atol=1e-9, fdm=_fdm, kwargs...)
     @test r_res !== f_res !== nothing  # Check the rule was defined
     r_fx, prop_rule = r_res
     f_fx, f_∂x = f_res
-    @testset "$f at $x, $(nameof(rule))" for (rule, fx, ∂x) in ((rrule, r_fx, prop_rule(1)), (frule, f_fx, f_∂x))
+    @testset "$f at $x, $(nameof(rule))" for (rule, fx, ∂x) in (
+        (rrule, r_fx, prop_rule(1)),
+        (frule, f_fx, f_∂x)
+    )
         @test fx == f(x)  # Check we still get the normal value, right
 
         if rule == rrule
             ∂self, ∂x = ∂x
             @test ∂self === NO_FIELDS
         end
-        @test isapprox(∂x, fdm(f, x);
-                       rtol=rtol, atol=atol, kwargs...)
+        @test isapprox(∂x, fdm(f, x); rtol=rtol, atol=atol, kwargs...)
     end
 end