value_and_gradient / value_and_pullback with kwargs (#482)

* Fix kwarg propagation

* Check the top-level interface works with kwargs

* Check all methods of value_and_gradient and value_and_pullback work with kwargs

* Bump patch version

Project.toml

@@ -1,7 +1,7 @@
 name = "Mooncake"
 uuid = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 authors = ["Will Tebbutt, Hong Ge, and contributors"]
-version = "0.4.86"
+version = "0.4.87"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/interpreter/s2s_reverse_mode_ad.jl

@@ -1008,21 +1008,22 @@ function Base.showerror(io::IO, err::MooncakeRuleCompilationError)
 end
 
 """
-    build_rrule(args...; debug_mode=false)
+    build_rrule(args...; kwargs...)
 
-Helper method. Only uses static information from `args`.
+Helper method: equivalent to extracting the signature from `args` and calling
+`build_rrule(sig; kwargs...)`.
 """
-function build_rrule(args...; debug_mode=false)
+function build_rrule(args...; kwargs...)
     interp = get_interpreter()
-    return build_rrule(interp, _typeof(TestUtils.__get_primals(args)); debug_mode)
+    return build_rrule(interp, _typeof(TestUtils.__get_primals(args)); kwargs...)
 end
 
 """
-    build_rrule(sig::Type{<:Tuple})
+    build_rrule(sig::Type{<:Tuple}; kwargs...)
 
-Equivalent to `build_rrule(Mooncake.get_interpreter(), sig)`.
+Helper method: Equivalent to `build_rrule(Mooncake.get_interpreter(), sig; kwargs...)`.
 """
-build_rrule(sig::Type{<:Tuple}) = build_rrule(get_interpreter(), sig)
+build_rrule(sig::Type{<:Tuple}; kwargs...) = build_rrule(get_interpreter(), sig; kwargs...)
 
 const MOONCAKE_INFERENCE_LOCK = ReentrantLock()
 

test/interface.jl

@@ -36,15 +36,16 @@ end
             ((x...) -> x[1] + x[2], randn(Float64), randn(Float64)),
             (sum, randn(10)),
         ]
-            rule = build_rrule(fargs...)
+            kwargs = (debug_mode=false, silence_debug_messages=true)
+            rule = build_rrule(fargs...; kwargs...)
             f, args... = fargs
             v, dfargs = value_and_gradient!!(rule, fargs...)
             @test v == f(args...)
             for (arg, darg) in zip(fargs, dfargs)
                 @test tangent_type(typeof(arg)) == typeof(darg)
             end
 
-            cache = Mooncake.prepare_gradient_cache(fargs...)
+            cache = Mooncake.prepare_gradient_cache(fargs...; kwargs...)
             _v, _dfargs = value_and_gradient!!(cache, fargs...)
             @test _v == v
             for (arg, darg) in zip(fargs, _dfargs)
@@ -74,15 +75,16 @@ end
             (randn(), sin, randn(Float64)),
             (randn(), sum, randn(Float64)),
         ]
-            rule = build_rrule(fargs...)
+            kwargs = (debug_mode=false, silence_debug_messages=true)
+            rule = build_rrule(fargs...; kwargs...)
             f, args... = fargs
             v, dfargs = value_and_pullback!!(rule, ȳ, fargs...)
             @test v == f(args...)
             for (arg, darg) in zip(fargs, dfargs)
                 @test tangent_type(typeof(arg)) == typeof(darg)
             end
 
-            cache = Mooncake.prepare_pullback_cache(fargs...)
+            cache = Mooncake.prepare_pullback_cache(fargs...; kwargs...)
             _v, _dfargs = value_and_pullback!!(cache, ȳ, fargs...)
             @test _v == v
             for (arg, darg) in zip(fargs, _dfargs)

test/interpreter/s2s_reverse_mode_ad.jl

@@ -349,4 +349,16 @@ end
         stmts = Mooncake.stmt(ir.stmts)
         @test !any(x -> Meta.isexpr(x, :new) && x.args[1] <: Base.RefValue, stmts)
     end
+    @testset "build_rrule methods all accept kwargs" begin
+        args = (sin, 5.0)
+        sig = typeof(args)
+        rule_sig = build_rrule(sig; debug_mode=false, silence_debug_messages=true)
+        @test rule_sig == rrule!!
+        rule_args = build_rrule(args...; debug_mode=false, silence_debug_messages=true)
+        @test rule_args == rrule!!
+        rule_debug_sig = build_rrule(sig; debug_mode=true, silence_debug_messages=true)
+        @test rule_debug_sig isa Mooncake.DebugRRule
+        rule_debug_args = build_rrule(args...; debug_mode=true, silence_debug_messages=true)
+        @test rule_debug_args == rule_debug_sig
+    end
 end