Merge pull request #58 from JuliaDiff/aa/checked

Add `_checked_rrule`, which errors if `rrule` returns `nothing`

Author: ararslan

src/rules.jl

@@ -377,6 +377,24 @@ See also: [`frule`](@ref), [`AbstractRule`](@ref), [`@scalar_rule`](@ref)
 """
 rrule(::Any, ::Vararg{Any}) = nothing
 
+@noinline function _throw_checked_rrule_error(f, args...; kwargs...)
+    io = IOBuffer()
+    print(io, "can't differentiate `", f, '(')
+    join(io, map(arg->string("::", typeof(arg)), args), ", ")
+    if !isempty(kwargs)
+        print(io, ";")
+        join(io, map(((k, v),)->string(k, "=", v), kwargs), ", ")
+    end
+    print(io, ")`; no matching `rrule` is defined")
+    throw(ArgumentError(String(take!(io))))
+end
+
+function _checked_rrule(f, args...; kwargs...)
+    r = rrule(f, args...; kwargs...)
+    r isa Nothing && _throw_checked_rrule_error(f, args...; kwargs...)
+    return r
+end
+
 #####
 ##### macros
 #####

src/rules/mapreduce.jl

@@ -7,12 +7,7 @@ function rrule(::typeof(map), f, xs...)
     ∂xs = ntuple(length(xs)) do i
         Rule() do ȳ
             map(ȳ, xs...) do ȳi, xis...
-                r = rrule(f, xis...)
-                if r === nothing
-                    throw(ArgumentError("can't differentiate `map` with `$f`; no `rrule` " *
-                                        "is defined for `$f$xis`"))
-                end
-                _, ∂xis = r
+                _, ∂xis = _checked_rrule(f, xis...)
                 extern(∂xis[i](ȳi))
             end
         end

test/rules.jl

@@ -1,4 +1,5 @@
 cool(x) = x + 1
+cool(x, y) = x + y + 1
 
 @testset "rules" begin
     @testset "frule and rrule" begin
@@ -45,5 +46,14 @@ cool(x) = x + 1
         @test ChainRules._update!(X, Y) == (A=[2 0; 0 2], B=[4 4; 4 4])
         @test X.A != Y.A
         @test X.B != Y.B
+
+        try
+            # We defined a 2-arg method for `cool` but no `rrule`
+            ChainRules._checked_rrule(cool, 1.0, 2.0)
+        catch e
+            @test e isa ArgumentError
+            @test e.msg == "can't differentiate `cool(::Float64, ::Float64)`; no " *
+                           "matching `rrule` is defined"
+        end
     end
 end