Update test/rules.jl

error ratehr than Assert
cleanup

Update src/rule_definition_tools.jl

Co-Authored-By: Nick Robinson <npr251@gmail.com>
Add more complex Wirtinger Scalar Rule Test

Author: oxinabox

src/differentials.jl

@@ -241,9 +241,10 @@ const NO_FIELDS = DNE()
 """
     differential(𝒟::Type, der)
 
-For some differential (e.g. a `Number`, `AbstractDifferential`, `Matrix`, etc.),
-convert it to another differential that is more suited for the domain given by
-the type 𝒟.
+Converts, if required, a differential object `der`
+(e.g. a `Number`, `AbstractDifferential`, `Matrix`, etc.),
+to another  differential that is more suited for the domain given by the type 𝒟.
+Often this will behave as the identity function on `der`.
 """
 function differential(::Type{<:Union{<:Real, AbstractArray{<:Real}}}, w::Wirtinger)
     return wirtinger_primal(w) + wirtinger_conjugate(w)

src/rule_definition_tools.jl

@@ -183,7 +183,7 @@ macro scalar_rule(call, maybe_setup, partials...)
         if Meta.isexpr(partial, :tuple)
             partial
         else
-            @assert length(inputs) == 1
+            length(inputs) == 1 || error("Invalid use of `@scalar_rule`")
             Expr(:tuple, partial)
         end
     end
@@ -192,7 +192,7 @@ macro scalar_rule(call, maybe_setup, partials...)
     # Main body: defining the results of the frule/rrule
 
     # An expression that when evaluated will return the type of the input domain.
-    # Multiple repetitions of this expression should optimize ot. But if it does not then
+    # Multiple repetitions of this expression should optimize out. But if it does not then
     # may need to move its definition into the body of the `rrule`/`frule`
     𝒟 = :(typeof(first(promote($(call.args[2:end]...)))))
 

test/rules.jl

@@ -29,44 +29,123 @@ _second(t) = Base.tuple_type_head(Base.tuple_type_tail(t))
                         Tuple{typeof(rrule),typeof(cool),String}])
     @test cool_methods == only_methods
 
-    frx, fr = frule(cool, 1)
+    frx, cool_pushforward = frule(cool, 1)
     @test frx == 2
-    @test fr(NamedTuple(), 1) == (1,)
-    rrx, (rr) = rrule(cool, 1)
-    self, rr1 = rr(1)
+    @test cool_pushforward(NamedTuple(), 1) == (1,)
+    rrx, cool_pullback = rrule(cool, 1)
+    self, rr1 = cool_pullback(1)
     @test self == NO_FIELDS
     @test rrx == 2
     @test rr1 == 1
 end
 
 
-@testset "Wirtinger scalar_rule" begin
+@testset "Basic Wirtinger scalar_rule" begin
     myabs2(x) = abs2(x)
     @scalar_rule(myabs2(x), Wirtinger(x', x))
 
-    # real input
-    x = rand(Float64)
-    f, pushforward = frule(myabs2, x)
-    @test f === x^2
+    @testset "real input" begin
+        # even though our rule was define in terms of Wirtinger,
+        # pushforward result will be real as real (even if seed is Compex)
 
-    df = @inferred pushforward(NamedTuple(), One())
-    @test df === (x + x,)
+        x = rand(Float64)
+        f, myabs2_pushforward = frule(myabs2, x)
+        @test f === x^2
 
+        Δ = One()
+        df = @inferred myabs2_pushforward(NamedTuple(), Δ)
+        @test df === (x + x,)
 
-    Δ = rand(Complex{Int64})
-    df = @inferred pushforward(NamedTuple(), Δ)
-    @test df === (Δ * (x + x),)
+        Δ = rand(Complex{Int64})
+        df = @inferred myabs2_pushforward(NamedTuple(), Δ)
+        @test df === (Δ * (x + x),)
+    end
 
+    @testset "complex input" begin
+        z = rand(Complex{Float64})
+        f, myabs2_pushforward = frule(myabs2, z)
+        @test f === abs2(z)
 
-    # complex input
-    z = rand(Complex{Float64})
-    f, pushforward = frule(myabs2, z)
-    @test f === abs2(z)
+        df = @inferred myabs2_pushforward(NamedTuple(), One())
+        @test df === (Wirtinger(z', z),)
+
+        Δ = rand(Complex{Int64})
+        df = @inferred myabs2_pushforward(NamedTuple(), Δ)
+        @test df === (Wirtinger(Δ * z', Δ * z),)
+    end
+end
 
-    df = @inferred pushforward(NamedTuple(), One())
-    @test df === (Wirtinger(z', z),)
 
-    Δ = rand(Complex{Int64})
-    df = @inferred pushforward(NamedTuple(), Δ)
-    @test df === (Wirtinger(Δ * z', Δ * z),)
+@testset "Advanced Wirtinger @scalar_rule: abs_to_pow" begin
+    # This is based on SimeonSchaub excellent example:
+    # https://gist.github.com/simeonschaub/a6dfcd71336d863b3777093b3b8d9c97
+
+    # This is much more complex than the previous case
+    # as it has many different types
+    # depending on input, and the output types do not always agree
+
+    abs_to_pow(x, p) = abs(x)^p
+    @scalar_rule(
+        abs_to_pow(x::Real, p),
+        (
+            p == 0 ? Zero() : p * abs_to_pow(x, p-1) * sign(x),
+            Ω * log(abs(x))
+        )
+    )
+
+    @scalar_rule(
+        abs_to_pow(x::Complex, p),
+        @setup(u = abs(x)),
+        (
+            p == 0 ? Zero() : p * u^(p-1) * Wirtinger(x' / 2u, x / 2u),
+            Ω * log(abs(x))
+        )
+    )
+
+
+    f = abs_to_pow
+    @testset "f($x, $p)" for (x, p) in Iterators.product(
+        (2, 3.4, -2.1, -10+0im, 2.3-2im),
+        (0, 1, 2, 4.3, -2.1, 1+.2im)
+    )
+        expected_type_df_dx =
+            if iszero(p)
+                Zero
+            elseif typeof(x) <: Complex
+                Wirtinger
+            elseif typeof(p) <: Complex
+                Complex
+            else
+                Real
+            end
+
+        expected_type_df_dp =
+            if typeof(p) <: Real
+                Real
+            else
+                Complex
+            end
+
+
+        res = frule(f, x, p)
+        @test res !== nothing  # Check the rule was defined
+        fx, f_pushforward = res
+        df(Δx, Δp) = f_pushforward(NamedTuple(), Δx, Δp)
+
+        df_dx, = df(One(), Zero())
+        df_dp,= df(Zero(), One())
+        @test fx == f(x, p)  # Check we still get the normal value, right
+        @test extern(df_dx) isa expected_type_df_dx
+        @test extern(df_dp) isa expected_type_df_dp
+
+
+        res = rrule(f, x, p)
+        @test res !== nothing  # Check the rule was defined
+        fx, f_pullback = res
+        dself, df_dx, df_dp = f_pullback(One())
+        @test fx == f(x, p)  # Check we still get the normal value, right
+        @test dself == NO_FIELDS
+        @test extern(df_dx) isa expected_type_df_dx
+        @test extern(df_dp) isa expected_type_df_dp
+    end
 end