Restrict kwarg type in code generated by @non_differentiable macro

oxinabox · oxinabox · commit eba7477e3a62 · 2024-01-29T12:53:44.000+08:00
diff --git a/src/rule_definition_tools.jl b/src/rule_definition_tools.jl
@@ -445,7 +445,7 @@ function _nondiff_rrule_expr(__source__, primal_sig_parts, primal_invoke)
     return @strip_linenos quote
         # Manually defined kw version to save compiler work. See explanation in rules.jl
         function (::Core.kwftype(typeof(rrule)))(
-            $(esc(kwargs))::Any, ::typeof(rrule), $(esc_primal_sig_parts...)
+            $(esc(kwargs))::NamedTuple, ::typeof(rrule), $(esc_primal_sig_parts...)
         )
             return ($(esc(_with_kwargs_expr(primal_invoke, kwargs))), $pullback_expr)
         end

Original file line number	Diff line number	Diff line change
`@@ -445,7 +445,7 @@ function _nondiff_rrule_expr(__source__, primal_sig_parts, primal_invoke)`
`445`	`445`	`return @strip_linenos quote`
`446`	`446`	`# Manually defined kw version to save compiler work. See explanation in rules.jl`
`447`	`447`	`function (::Core.kwftype(typeof(rrule)))(`
`448`		`- $(esc(kwargs))::Any, ::typeof(rrule), $(esc_primal_sig_parts...)`
	`448`	`+ $(esc(kwargs))::NamedTuple, ::typeof(rrule), $(esc_primal_sig_parts...)`
`449`	`449`	`)`
`450`	`450`	`return ($(esc(_with_kwargs_expr(primal_invoke, kwargs))), $pullback_expr)`
`451`	`451`	`end`