Merge pull request #217 from JuliaDiff/ox/nondiff_kw

Add kwarg support to at-nondifferentiable

Author: oxinabox

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRulesCore"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "0.9.8"
+version = "0.9.9"
 
 [deps]
 MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"

src/rule_definition_tools.jl

@@ -295,7 +295,8 @@ macro non_differentiable(sig_expr)
     primal_sig_parts = [:(::typeof($primal_name)), constrained_args...]
 
     unconstrained_args = _unconstrain.(constrained_args)
-    primal_invoke = Expr(:call, esc(primal_name), esc.(unconstrained_args)...)
+
+    primal_invoke = :($(primal_name)($(unconstrained_args...); kwargs...))
 
     quote
         $(_nondiff_frule_expr(primal_sig_parts, primal_invoke))
@@ -304,28 +305,27 @@ macro non_differentiable(sig_expr)
 end
 
 function _nondiff_frule_expr(primal_sig_parts, primal_invoke)
-    return Expr(
-        :(=),
-        Expr(:call, :(ChainRulesCore.frule), esc(:_), esc.(primal_sig_parts)...),
-        # Julia functions always only have 1 output, so just return a single DoesNotExist()
-        Expr(:tuple, primal_invoke, DoesNotExist()),
-    )
+    return esc(:(
+        function ChainRulesCore.frule($(gensym(:_)), $(primal_sig_parts...); kwargs...)
+            # Julia functions always only have 1 output, so return a single DoesNotExist()
+            return ($primal_invoke, DoesNotExist())
+        end
+    ))
 end
 
 function _nondiff_rrule_expr(primal_sig_parts, primal_invoke)
     num_primal_inputs = length(primal_sig_parts) - 1
     primal_name = first(primal_invoke.args)
     pullback_expr = Expr(
         :function,
-        Expr(:call, esc(propagator_name(primal_name, :pullback)), esc(:_)),
+        Expr(:call, propagator_name(primal_name, :pullback), :_),
         Expr(:tuple, NO_FIELDS, ntuple(_->DoesNotExist(), num_primal_inputs)...)
     )
-    rrule_defn = Expr(
-        :(=),
-        Expr(:call, :(ChainRulesCore.rrule), esc.(primal_sig_parts)...),
-        Expr(:tuple, primal_invoke, pullback_expr),
-    )
-    return rrule_defn
+    return esc(:(
+        function ChainRulesCore.rrule($(primal_sig_parts...); kwargs...)
+            return ($primal_invoke, $pullback_expr)
+        end
+    ))
 end
 
 

test/rule_definition_tools.jl

@@ -73,6 +73,31 @@ end
             @test rrule(pointy_identity, 2.0) == nothing
         end
 
+        @testset "kwargs" begin
+            kw_demo(x; kw=2.0) = x + kw
+            @non_differentiable kw_demo(::Any)
+
+            @testset "not setting kw" begin
+                @assert kw_demo(1.5) == 3.5
+
+                res, pullback = rrule(kw_demo, 1.5)
+                @test res == 3.5
+                @test pullback(4.1) == (NO_FIELDS, DoesNotExist())
+
+                @test frule((Zero(), 11.1), kw_demo, 1.5) == (3.5, DoesNotExist())
+            end
+
+            @testset "setting kw" begin
+                @assert kw_demo(1.5; kw=3.0) == 4.5
+
+                res, pullback = rrule(kw_demo, 1.5; kw=3.0)
+                @test res == 4.5
+                @test pullback(1.1) == (NO_FIELDS, DoesNotExist())
+
+                @test frule((Zero(), 11.1), kw_demo, 1.5; kw=3.0) == (4.5, DoesNotExist())
+            end
+        end
+
         @testset "Not supported (Yet)" begin
             # Varargs are not supported
             @test_macro_throws ErrorException @non_differentiable vararg1(xs...)