Merge pull request #651 from JuliaDiff/ox/non_diff_ambig

Fix ambig with configured frule created by non_differentiable rule
JuliaDiff · Jan 12, 2024 · 77748b3 · 77748b3 · oxinabox · Jan 12, 2024
2 parents 17a83d4 + 161e5f1
commit 77748b3
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 3 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "ChainRulesCore"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "1.19.0"
+version = "1.19.1"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"

diff --git a/src/rule_definition_tools.jl b/src/rule_definition_tools.jl
@@ -403,13 +403,13 @@ function _nondiff_frule_expr(__source__, primal_sig_parts, primal_invoke)
         function (::Core.kwftype(typeof(ChainRulesCore.frule)))(
             @nospecialize($kwargs::Any),
             frule::typeof(ChainRulesCore.frule),
-            @nospecialize(::Any),
+            @nospecialize(::Tuple),
             $(map(esc, primal_sig_parts)...),
         )
             return ($(esc(_with_kwargs_expr(primal_invoke, kwargs))), NoTangent())
         end
         function ChainRulesCore.frule(
-            @nospecialize(::Any), $(map(esc, primal_sig_parts)...)
+            @nospecialize(::Tuple), $(map(esc, primal_sig_parts)...)
         )
             $(__source__)
             # Julia functions always only have 1 output, so return a single NoTangent()

diff --git a/test/rule_definition_tools.jl b/test/rule_definition_tools.jl
@@ -215,6 +215,24 @@ end
             @test pullback(4.5) == (NoTangent(), NoTangent(), NoTangent())
         end
 
+        @testset "interactions with configs" begin
+            struct AllConfig <: RuleConfig{Union{HasForwardsMode,NoReverseMode}} end
+
+            foo_ndc1(x) = string(x)
+            @non_differentiable foo_ndc1(x)
+            @test frule(AllConfig(), (NoTangent(), NoTangent()), foo_ndc1, 2.0) == (string(2.0), NoTangent())
+            r1, pb1 = rrule(AllConfig(), foo_ndc1, 2.0)
+            @test r1 == string(2.0)
+            @test pb1(NoTangent()) == (NoTangent(), NoTangent())
+
+            foo_ndc2(x; y=0) = string(x + y)
+            @non_differentiable foo_ndc2(x)
+            @test frule(AllConfig(), (NoTangent(), NoTangent()), foo_ndc2, 2.0; y=4.0) == (string(6.0), NoTangent())
+            r2, pb2 = rrule(AllConfig(), foo_ndc2, 2.0; y=4.0)
+            @test r2 == string(6.0)
+            @test pb2(NoTangent()) == (NoTangent(), NoTangent())
+        end
+
         @testset "Not supported (Yet)" begin
             # Where clauses are not supported.
             @test_macro_throws(