Support kwargs in rrule_via_ad (#1055)

* allow kwargs

* add test for kwarg support

* bump patch

* unwrap closure

* first solution

* second solution kwf()

* Revert "second solution kwf()"

This reverts commit b53a381.

* avoid creating closure unnecesasrily

* short function

* first instead of only

Project.toml

@@ -1,6 +1,6 @@
 name = "Zygote"
 uuid = "e88e6eb3-aa80-5325-afca-941959d7151f"
-version = "0.6.19"
+version = "0.6.20"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

src/compiler/chainrules.jl

@@ -174,10 +174,18 @@ As per [`chain_rrule`](@ref) but with support for kwargs.
   return y, kw_zpullback
 end
 
+function ChainRulesCore.rrule_via_ad(config::ZygoteRuleConfig, f_args...; kwargs...)
+    # create a closure to work around _pullback not accepting kwargs
+    # but avoid creating a closure unnecessarily (pullbacks of closures do not infer)
+    y, pb = if !isempty(kwargs)
+        kwf() = first(f_args)(Base.tail(f_args)...; kwargs...)
+        _y, _pb = _pullback(config.context, kwf)
+        _y, Δ -> first(_pb(Δ)).f_args  # `first` should be `only`
+    else
+        _pullback(config.context, f_args...)
+    end
 
-function ChainRulesCore.rrule_via_ad(config::ZygoteRuleConfig, f, args...)
-    y, pb = _pullback(config.context, f, args...)
-    ad_pullback(Δ) = zygote2differential(pb(wrap_chainrules_output(Δ)), (f, args...))
+    ad_pullback(Δ) = zygote2differential(pb(wrap_chainrules_output(Δ)), f_args)
     return y, ad_pullback
 end
 

test/chainrules.jl

@@ -275,6 +275,13 @@ end
         test_rrule(ZygoteRuleConfig(), getindex, rand(5), 3; rrule_f=rrule_via_ad)
     end
 
+    @testset "kwargs" begin
+        test_rrule(
+            ZygoteRuleConfig(), sum, [1.0 2; 3 4];
+            rrule_f=rrule_via_ad, check_inferred=false, fkwargs=(;dims=1)
+        )
+    end
+
     @testset "struct" begin
         struct Foo
             x