Correct chainrules for abs2, abs, conj and angle

src/rulesets/Base/fastmath_able.jl

@@ -36,10 +36,88 @@ let
 
 
         # Unary complex functions
-        @scalar_rule abs(x::Real) sign(x)
-        @scalar_rule abs2(x) 2x
-        @scalar_rule angle(x::Real) Zero()
-        @scalar_rule conj(x::Real) One()
+        ## abs
+        function frule((_, Δx), ::typeof(abs), x::Real)
+            return abs(x), sign(x) * real(Δx)
+        end
+        function frule((_, Δz), ::typeof(abs), z::Complex)
+            Ω = abs(z)
+            return Ω, (real(z) * real(Δz) + imag(z) * imag(Δz)) / Ω
+        end
+
+        function rrule(::typeof(abs), x::Real)
+            function abs_pullback(ΔΩ)
+                return (NO_FIELDS, real(ΔΩ)*sign(x))
+            end
+            return abs(x), abs_pullback
+        end
+        function rrule(::typeof(abs), z::Complex)
+            Ω = abs(z)
+            function abs_pullback(ΔΩ)
+                Δu = real(ΔΩ) 
+                return (NO_FIELDS, Δu*z/Ω)
+            end
+            return Ω, abs_pullback
+        end
+
+        ## abs2
+        function frule((_, Δx), ::typeof(abs2), x::Real)
+            return abs2(x), 2x * real(Δx)
+        end
+        function frule((_, Δz), ::typeof(abs2), z::Complex)
+            return abs2(z), 2 * (real(z) * real(Δz) + imag(z) * imag(Δz))
+        end
+
+        function rrule(::typeof(abs2), x::Real)
+            function abs2_pullback(Δx)
+                return (NO_FIELDS, 2real(Δx)*x)
+            end
+            return abs2(x), abs2_pullback
+        end
+        function rrule(::typeof(abs2), z::Complex)
+            function abs2_pullback(Δf)
+                Δu = real(Δf)
+                return (NO_FIELDS, 2real(Δu)*z)
+            end
+            return abs2(z), abs2_pullback
+        end
+
+        ## conj
+        function frule((_, Δz), ::typeof(conj), z::Union{Real, Complex})
+            return conj(z), conj(Δz) 
+        end
+        function rrule(::typeof(conj), z::Union{Real, Complex})
+            function conj_pullback(Δf)
+                return (NO_FIELDS, conj(Δf))
+            end
+            return conj(z), conj_pullback
+        end
+
+        ## angle
+        function frule((_, Δz), ::typeof(angle), x::Real)
+            Δx, Δy = reim(Δz)
+            return angle(x), Δy/x  
+        end
+        function frule((_, Δz), ::typeof(angle), z::Complex)
+            x,  y  = reim(z)
+            Δx, Δy = reim(Δz)
+            return angle(z), (-y*Δx + x*Δy)/abs2(z)  
+        end
+        function rrule(::typeof(angle), x::Real)
+            function angle_pullback(Δf)
+                Δu, Δv = reim(Δf)
+                return (NO_FIELDS, im*Δu/x)
+            end
+            return angle(x), angle_pullback 
+        end
+        function rrule(::typeof(angle), z::Complex)
+            function angle_pullback(Δf)
+                x,  y  = reim(z)
+                Δu, Δv = reim(Δf)
+                return (NO_FIELDS, (-y + im*x)*Δu/abs2(z))
+            end
+            return angle(z), angle_pullback 
+        end
 
         # Binary functions
         @scalar_rule hypot(x::Real, y::Real) (x / Ω, y / Ω)

test/rulesets/Base/fastmath_able.jl

@@ -1,4 +1,42 @@
 # Add tests to the quote for functions with  FastMath varients.
+function jacobian_via_frule(f,z)
+    du_dx, dv_dx = reim(frule((Zero(), 1),f,z)[2])
+    du_dy, dv_dy = reim(frule((Zero(),im),f,z)[2])
+    return [
+        du_dx  du_dy
+        dv_dx  dv_dy
+    ]
+end
+function jacobian_via_rrule(f,z)
+    _, pullback = rrule(f,z)
+    du_dx, du_dy = reim(pullback( 1)[2])
+    dv_dx, dv_dy = reim(pullback(im)[2])
+    return [
+        du_dx  du_dy
+        dv_dx  dv_dy
+    ]
+end
+
+function jacobian_via_fdm(f, z::Union{Real, Complex})
+    fR2((x, y)) = (collect ∘ reim ∘ f)(x + im*y)
+    v = float([real(z)
+               imag(z)])
+    j = jacobian(central_fdm(5,1), fR2, v)[1]
+    if size(j) == (2,2)
+        j
+    elseif size(j) == (1, 2)
+        [j 
+         false false]
+    else
+        error("Invalid Jacobian size $(size(j))")
+    end
+end
+
+function complex_jacobian_test(f, z)
+    @test jacobian_via_fdm(f, z) ≈ jacobian_via_frule(f, z) 
+    @test jacobian_via_fdm(f, z) ≈ jacobian_via_rrule(f, z)
+end
+
 const FASTABLE_AST = quote
     @testset "Trig" begin
         @testset "Basics" for x = (Float64(π)-0.01, Complex(π, π/2))
@@ -47,13 +85,12 @@ const FASTABLE_AST = quote
             end
         end
     end
-
+    
     @testset "Unary complex functions" begin
-        for x in (-4.1, 6.4)
-            test_scalar(abs, x)
-            test_scalar(angle, x)
-            test_scalar(abs2, x)
-            test_scalar(conj, x)
+        for f ∈ (abs, abs2, angle, conj), z ∈ (-4.1-0.02im, 6.4, 3 + im)
+            @testset "Unary complex functions f = $f, z = $z" begin
+                complex_jacobian_test(f, z)
+            end
         end
     end
 
@@ -73,8 +110,6 @@ const FASTABLE_AST = quote
         rrule_test(f, Δz, (x, x̄), (y, ȳ))
     end
 
-
-
     @testset "sign" begin
         @testset "at points" for x in (-1.1, -1.1, 0.5, 100)
             test_scalar(sign, x)