make @non_differentiable use identical pullbacks when possible

Fixes #678
JuliaDiff · May 31, 2024 · 191eb47 · 191eb47
1 parent fa530b9
commit 191eb47
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 3 deletions.
diff --git a/src/rule_definition_tools.jl b/src/rule_definition_tools.jl
@@ -418,6 +418,14 @@ function _nondiff_frule_expr(__source__, primal_sig_parts, primal_invoke)
     end
 end
 
+struct NonDiffPullback{T<:Tuple{Vararg{NoTangent}}} <: Function
+    v::T
+end
+
+function (@nospecialize pb::NonDiffPullback)(@nospecialize ::Any)
+    return pb.v
+end
+
 function tuple_expression(primal_sig_parts)
     has_vararg = _isvararg(primal_sig_parts[end])
     return if !has_vararg
@@ -436,9 +444,7 @@ function _nondiff_rrule_expr(__source__, primal_sig_parts, primal_invoke)
     tup_expr = tuple_expression(primal_sig_parts)
     primal_name = first(primal_invoke.args)
     pullback_expr = @strip_linenos quote
-        function $(esc(propagator_name(primal_name, :pullback)))(@nospecialize(_))
-            return $(tup_expr)
-        end
+        NonDiffPullback($(tup_expr))
     end
 
     @gensym kwargs

diff --git a/test/rule_definition_tools.jl b/test/rule_definition_tools.jl
@@ -42,6 +42,32 @@ end
 
 @testset "rule_definition_tools.jl" begin
     @testset "@non_differentiable" begin
+        @testset "`NonDiffPullback`" begin
+            NDP = ChainRulesCore.NonDiffPullback
+            for i in 0:5
+                tup = ntuple((_ -> NoTangent()), i)
+                ndp = NDP(tup)
+                @test ndp === @inferred NDP(tup)
+                @test tup === @inferred ndp(:arbitrary)
+                @test_throws MethodError ndp()
+                @test_throws MethodError ndp(1, 2)
+            end
+        end
+
+        @testset "issue #678: identical pullback objects" begin
+            issue_678_f(::Any) = nothing
+            issue_678_g(::Any) = nothing
+            issue_678_h(::Any...) = nothing
+            @non_differentiable issue_678_f(::Any)
+            @non_differentiable issue_678_g(::Any)
+            @non_differentiable issue_678_h(::Any...)
+            @test (
+                last(rrule(issue_678_f, 0.1)) ===
+                last(rrule(issue_678_g, 0.2)) ===
+                last(rrule(issue_678_h, 0.3))
+            )
+        end
+
         @testset "two input one output function" begin
             nondiff_2_1(x, y) = fill(7.5, 100)[x + y]
             @non_differentiable nondiff_2_1(::Any, ::Any)