Add missing complex tests and rules

Project.toml

@@ -11,7 +11,7 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
 ChainRulesCore = "0.9"
-ChainRulesTestUtils = "0.4.1"
+ChainRulesTestUtils = "0.4.2"
 Compat = "3"
 FiniteDifferences = "0.10"
 Reexport = "0.2"

src/ChainRules.jl

@@ -21,6 +21,7 @@ if VERSION < v"1.3.0-DEV.142"
     import LinearAlgebra: dot
 end
 
+include("rulesets/Base/utils.jl")
 include("rulesets/Base/base.jl")
 include("rulesets/Base/fastmath_able.jl")
 include("rulesets/Base/array.jl")

src/rulesets/Base/base.jl

@@ -3,15 +3,80 @@
 
 @scalar_rule one(x) zero(x)
 @scalar_rule zero(x) zero(x)
-@scalar_rule adjoint(x::Real) One()
 @scalar_rule transpose(x) One()
+
+# `adjoint`
+
+frule((_, Δz), ::typeof(adjoint), z::Number) = (z', Δz')
+
+function rrule(::typeof(adjoint), z::Number)
+    adjoint_pullback(ΔΩ) = (NO_FIELDS, ΔΩ')
+    return (z', adjoint_pullback)
+end
+
+# `real`
+
+@scalar_rule real(x::Real) One()
+
+frule((_, Δz), ::typeof(real), z::Number) = (real(z), real(Δz))
+
+function rrule(::typeof(real), z::Number)
+    # add zero(z) to embed the real number in the same number type as z
+    real_pullback(ΔΩ) = (NO_FIELDS, real(ΔΩ) + zero(z))
+    return (real(z), real_pullback)
+end
+
+# `imag`
+
 @scalar_rule imag(x::Real) Zero()
+
+frule((_, Δz), ::typeof(imag), z::Complex) = (imag(z), imag(Δz))
+
+function rrule(::typeof(imag), z::Complex)
+    imag_pullback(ΔΩ) = (NO_FIELDS, real(ΔΩ) * im)
+    return (imag(z), imag_pullback)
+end
+
+# `Complex`
+
+frule((_, Δz), ::Type{T}, z::Number) where {T<:Complex} = (T(z), Complex(Δz))
+function frule((_, Δx, Δy), ::Type{T}, x::Number, y::Number) where {T<:Complex}
+    return (T(x, y), Complex(Δx, Δy))
+end
+
+function rrule(::Type{T}, z::Complex) where {T<:Complex}
+    Complex_pullback(ΔΩ) = (NO_FIELDS, Complex(ΔΩ))
+    return (T(z), Complex_pullback)
+end
+function rrule(::Type{T}, x::Real) where {T<:Complex}
+    Complex_pullback(ΔΩ) = (NO_FIELDS, real(ΔΩ))
+    return (T(x), Complex_pullback)
+end
+function rrule(::Type{T}, x::Number, y::Number) where {T<:Complex}
+    Complex_pullback(ΔΩ) = (NO_FIELDS, real(ΔΩ), imag(ΔΩ))
+    return (T(x, y), Complex_pullback)
+end
+
+# `hypot`
+
 @scalar_rule hypot(x::Real) sign(x)
 
+function frule((_, Δz), ::typeof(hypot), z::Complex)
+    Ω = hypot(z)
+    ∂Ω = _realconjtimes(z, Δz) / ifelse(iszero(Ω), one(Ω), Ω)
+    return Ω, ∂Ω
+end
+
+function rrule(::typeof(hypot), z::Complex)
+    Ω = hypot(z)
+    function hypot_pullback(ΔΩ)
+        return (NO_FIELDS, (real(ΔΩ) / ifelse(iszero(Ω), one(Ω), Ω)) * z)
+    end
+    return (Ω, hypot_pullback)
+end
 
 @scalar_rule fma(x, y, z) (y, x, One())
 @scalar_rule muladd(x, y, z) (y, x, One())
-@scalar_rule real(x::Real) One()
 @scalar_rule rem2pi(x, r::RoundingMode) (One(), DoesNotExist())
 @scalar_rule(
     mod(x, y),

src/rulesets/Base/fastmath_able.jl

@@ -25,7 +25,7 @@ let
         @scalar_rule cbrt(x) inv(3 * Ω ^ 2)
         @scalar_rule inv(x) -(Ω ^ 2)
         @scalar_rule sqrt(x) inv(2Ω)
-        @scalar_rule exp(x::Real) Ω
+        @scalar_rule exp(x) Ω
         @scalar_rule exp10(x) Ω * log(oftype(x, 10))
         @scalar_rule exp2(x) Ω * log(oftype(x, 2))
         @scalar_rule expm1(x) exp(x)
@@ -37,37 +37,25 @@ let
 
         # Unary complex functions
         ## 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)) / ifelse(iszero(z), one(Ω), Ω)
+        function frule((_, Δx), ::typeof(abs), x::Union{Real, Complex})
+            Ω = abs(x)
+            signx = x isa Real ? sign(x) : x / ifelse(iszero(x), one(Ω), Ω)
             # `ifelse` is applied only to denominator to ensure type-stability.
+            return Ω, _realconjtimes(signx, Δx)
         end
 
-        function rrule(::typeof(abs), x::Real)
+        function rrule(::typeof(abs), x::Union{Real, Complex})
+            Ω = abs(x)
             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/ifelse(iszero(z), one(Ω), Ω))
-                # `ifelse` is applied only to denominator to ensure type-stability.
+                signx = x isa Real ? sign(x) : x / ifelse(iszero(x), one(Ω), Ω)
+                return (NO_FIELDS, signx * real(ΔΩ))
             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))
+        function frule((_, Δz), ::typeof(abs2), z::Union{Real, Complex})
+            return abs2(z), 2 * _realconjtimes(z, Δz)
         end
 
         function rrule(::typeof(abs2), z::Union{Real, Complex})
@@ -90,20 +78,13 @@ let
         end
 
         ## angle
-        function frule((_, Δz), ::typeof(angle), x::Real)
-            Δx, Δy = reim(Δz)
-            return angle(x), Δy/ifelse(iszero(x), one(x), x)
-            # `ifelse` is applied only to denominator to ensure type-stability.
-        end
-        function frule((_, Δz)::Tuple{<:Any, <:Real}, ::typeof(angle), x::Real)
-            return angle(x), Zero()
-        end
-        function frule((_, Δz), ::typeof(angle), z::Complex)
-            x,  y  = reim(z)
-            Δx, Δy = reim(Δz)
-            return angle(z), (-y*Δx + x*Δy)/ifelse(iszero(z), one(z), abs2(z))
+        function frule((_, Δx), ::typeof(angle), x)
+            Ω = angle(x)
             # `ifelse` is applied only to denominator to ensure type-stability.
+            ∂Ω = _imagconjtimes(x, Δx) / ifelse(iszero(x), one(x), abs2(x))
+            return Ω, ∂Ω
         end
+
         function rrule(::typeof(angle), x::Real)
             function angle_pullback(ΔΩ::Real)
                 return (NO_FIELDS, Zero())
@@ -126,7 +107,30 @@ let
         end
 
         # Binary functions
-        @scalar_rule hypot(x::Real, y::Real) (x / Ω, y / Ω)
+
+        # `hypot`
+
+        function frule(
+            (_, Δx, Δy),
+            ::typeof(hypot),
+            x::T,
+            y::T,
+        ) where {T<:Union{Real,Complex}}
+            Ω = hypot(x, y)
+            n = ifelse(iszero(Ω), one(Ω), Ω)
+            ∂Ω = (_realconjtimes(x, Δx) + _realconjtimes(y, Δy)) / n
+            return Ω, ∂Ω
+        end
+
+        function rrule(::typeof(hypot), x::T, y::T) where {T<:Union{Real,Complex}}
+            Ω = hypot(x, y)
+            function hypot_pullback(ΔΩ)
+                c = real(ΔΩ) / ifelse(iszero(Ω), one(Ω), Ω)
+                return (NO_FIELDS, @thunk(c * x), @thunk(c * y))
+            end
+            return (Ω, hypot_pullback)
+        end
+
         @scalar_rule x + y (One(), One())
         @scalar_rule x - y (One(), -1)
         @scalar_rule x / y (inv(y), -((x / y) / y))
@@ -146,30 +150,21 @@ let
         @scalar_rule +x One()
         @scalar_rule -x -1
 
-        function frule((_, Δx), ::typeof(sign), x::Real)
-            Ω = sign(x)
-            ∂Ω = _sign_jvp(Ω, x, Δx)
-            return Ω, ∂Ω
-        end
-        function frule((_, Δz), ::typeof(sign), z::Complex)
-            absz = abs(ifelse(iszero(z), one(z), z))
-            Ω = z / absz
-            ∂Ω = _sign_jvp(Ω, absz, Δz)
+        # `sign`
+
+        function frule((_, Δx), ::typeof(sign), x)
+            n = ifelse(iszero(x), one(x), abs(x))
+            Ω = x isa Real ? sign(x) : x / n
+            ∂Ω = Ω * (_imagconjtimes(Ω, Δx) / n) * im
             return Ω, ∂Ω
         end
 
-        function rrule(::typeof(sign), x::Real)
-            Ω = sign(x)
-            function sign_pullback(ΔΩ)
-                return (NO_FIELDS, _sign_jvp(Ω, x, ΔΩ))
-            end
-            return Ω, sign_pullback
-        end
-        function rrule(::typeof(sign), z::Complex)
-            absz = abs(ifelse(iszero(z), one(z), z))
-            Ω = z / absz
+        function rrule(::typeof(sign), x)
+            n = ifelse(iszero(x), one(x), abs(x))
+            Ω = x isa Real ? sign(x) : x / n
             function sign_pullback(ΔΩ)
-                return (NO_FIELDS, _sign_jvp(Ω, absz, ΔΩ))
+                ∂x = Ω * (_imagconjtimes(Ω, ΔΩ) / n) * im
+                return (NO_FIELDS, ∂x)
             end
             return Ω, sign_pullback
         end
@@ -196,9 +191,3 @@ let
     eval(fastable_ast)  # Get original definitions
     # we do this second so it overwrites anything we included by mistake in the fastable
 end
-
-# the jacobian for `sign` is symmetric; `_sign_jvp` gives both J * Δz and Jᵀ * ΔΩ for
-# output Ω, (co)tangent Δ, and real input x or the absolute value of complex input z
-_sign_jvp(Ω, absz, Δ) = Ω * ((imag(Δ) * real(Ω) - real(Δ) * imag(Ω)) / absz)im
-_sign_jvp(Ω::Real, x::Real, Δ) = (imag(Δ) * Ω / ifelse(iszero(x), one(x), x)) * im
-_sign_jvp(Ω::Real, x::Real, Δ::Real) = Zero()

src/rulesets/Base/utils.jl

@@ -0,0 +1,15 @@
+# real(conj(x) * y) avoiding computing the imaginary part if possible
+@inline _realconjtimes(x, y) = real(conj(x) * y)
+@inline _realconjtimes(x::Complex, y::Complex) = muladd(real(x), real(y), imag(x) * imag(y))
+@inline _realconjtimes(x::Real, y::Complex) = x * real(y)
+@inline _realconjtimes(x::Complex, y::Real) = real(x) * y
+@inline _realconjtimes(x::Real, y::Real) = x * y
+
+# imag(conj(x) * y) avoiding computing the real part if possible
+@inline _imagconjtimes(x, y) = imag(conj(x) * y)
+@inline function _imagconjtimes(x::Complex, y::Complex)
+    return muladd(-imag(x), real(y), real(x) * imag(y))
+end
+@inline _imagconjtimes(x::Real, y::Complex) = x * imag(y)
+@inline _imagconjtimes(x::Complex, y::Real) = -imag(x) * y
+@inline _imagconjtimes(x::Real, y::Real) = Zero()