Ensure scalar_rule rules are Number typed (#62)

Currently, `@scalar_rule` generates `rrule` methods with no type
constraints. This poses a problem for rules with different definitions
that have not yet been implemented. For example, the rule for matrix
exponential is quite different from that for scalar exponential, but in
the absence of an `rrule(::typeof(exp), ::AbstractMatrix)` method, the
incorrect fallback generated by `@scalar_rule` is used.

To solve this, `@scalar_rule` now allows type constaints, e.g.
`@scalar_rule(f(x::Complex), g(x))`, and it adds explicit `::Number`
constraints to the generated methods if no such constraints are
provided.

Author: ararslan

src/rules.jl

@@ -410,22 +410,29 @@ A convenience macro that generates simple scalar forward or reverse rules using
 the provided partial derivatives. Specifically, generates the corresponding
 methods for `frule` and `rrule`:
 
-    function ChainRules.frule(::typeof(f), x₁, x₂, ...)
+    function ChainRules.frule(::typeof(f), x₁::Number, x₂::Number, ...)
         Ω = f(x₁, x₂, ...)
         \$(statement₁, statement₂, ...)
         return Ω, (Rule((Δx₁, Δx₂, ...) -> ∂f₁_∂x₁ * Δx₁ + ∂f₁_∂x₂ * Δx₂ + ...),
                    Rule((Δx₁, Δx₂, ...) -> ∂f₂_∂x₁ * Δx₁ + ∂f₂_∂x₂ * Δx₂ + ...),
                    ...)
     end
 
-    function ChainRules.rrule(::typeof(f), x₁, x₂, ...)
+    function ChainRules.rrule(::typeof(f), x₁::Number, x₂::Number, ...)
         Ω = f(x₁, x₂, ...)
         \$(statement₁, statement₂, ...)
         return Ω, (Rule((ΔΩ₁, ΔΩ₂, ...) -> ∂f₁_∂x₁ * ΔΩ₁ + ∂f₂_∂x₁ * ΔΩ₂ + ...),
                    Rule((ΔΩ₁, ΔΩ₂, ...) -> ∂f₁_∂x₂ * ΔΩ₁ + ∂f₂_∂x₂ * ΔΩ₂ + ...),
                    ...)
     end
 
+If no type constraints in `f(x₁, x₂, ...)` within the call to `@scalar_rule` are
+provided, each parameter in the resulting `frule`/`rrule` definition is given a
+type constraint of `Number`.
+Constraints may also be explicitly be provided to override the `Number` constraint,
+e.g. `f(x₁::Complex, x₂)`, which will constrain `x₁` to `Complex` and `x₂` to
+`Number`.
+
 Note that the result of `f(x₁, x₂, ...)` is automatically bound to `Ω`. This
 allows the primal result to be conveniently referenced (as `Ω`) within the
 derivative/setup expressions.
@@ -458,7 +465,19 @@ macro scalar_rule(call, maybe_setup, partials...)
         partials = (maybe_setup, partials...)
     end
     @assert Meta.isexpr(call, :call)
-    f, inputs = esc(call.args[1]), esc.(call.args[2:end])
+    f = esc(call.args[1])
+    # Annotate all arguments in the signature as scalars
+    inputs = map(call.args[2:end]) do arg
+        esc(Meta.isexpr(arg, :(::)) ? arg : Expr(:(::), arg, :Number))
+    end
+    # Remove annotations and escape names for the call
+    for (i, arg) in enumerate(call.args)
+        if Meta.isexpr(arg, :(::))
+            call.args[i] = esc(first(arg.args))
+        else
+            call.args[i] = esc(arg)
+        end
+    end
     if all(Meta.isexpr(partial, :tuple) for partial in partials)
         forward_rules = Any[rule_from_partials(partial.args...) for partial in partials]
         reverse_rules = Any[]

src/rules/base.jl

@@ -38,7 +38,7 @@
 @scalar_rule(adjoint(x), Wirtinger(Zero(), One()))
 @scalar_rule(transpose(x), One())
 @scalar_rule(abs(x), sign(x))
-@scalar_rule(rem2pi(x, r), (One(), DNE()))
+@scalar_rule(rem2pi(x, r::RoundingMode), (One(), DNE()))
 @scalar_rule(+(x), One())
 @scalar_rule(-(x), -1)
 @scalar_rule(+(x, y), (One(), One()))

test/rules.jl

@@ -1,13 +1,25 @@
 cool(x) = x + 1
 cool(x, y) = x + y + 1
 
+_second(t) = Base.tuple_type_head(Base.tuple_type_tail(t))
+
 @testset "rules" begin
     @testset "frule and rrule" begin
         @test frule(cool, 1) === nothing
         @test frule(cool, 1; iscool=true) === nothing
         @test rrule(cool, 1) === nothing
         @test rrule(cool, 1; iscool=true) === nothing
+
         ChainRules.@scalar_rule(Main.cool(x), one(x))
+        @test hasmethod(rrule, Tuple{typeof(cool),Number})
+        ChainRules.@scalar_rule(Main.cool(x::String), "wow such dfdx")
+        @test hasmethod(rrule, Tuple{typeof(cool),String})
+        # Ensure those are the *only* methods that have been defined
+        cool_methods = Set(m.sig for m in methods(rrule) if _second(m.sig) == typeof(cool))
+        only_methods = Set([Tuple{typeof(rrule),typeof(cool),Number},
+                            Tuple{typeof(rrule),typeof(cool),String}])
+        @test cool_methods == only_methods
+
         frx, fr = frule(cool, 1)
         @test frx == 2
         @test fr(1) == 1