address comment

Author: nsajko

src/rule_definition_tools.jl

@@ -418,11 +418,10 @@ function _nondiff_frule_expr(__source__, primal_sig_parts, primal_invoke)
     end
 end
 
-function _make_pullback_for_non_differentiable(::Val{N}) where {N}
-    Vararg{Any,N}  # throw early for invalid `N`, must be nonnegative `Int`
-    function pullback_for_non_differentiable(::Any)
-        return ntuple(Returns(NoTangent()), Val(N))
-    end
+struct NonDiffPullback{N} <: Function end
+
+function (pb::NonDiffPullback{N})(::Any) where {N}
+    return ntuple(Returns(NoTangent()), Val(N))
 end
 
 function tuple_length_expression(primal_sig_parts)
@@ -443,7 +442,7 @@ function _nondiff_rrule_expr(__source__, primal_sig_parts, primal_invoke)
     tup_len_expr = tuple_length_expression(primal_sig_parts)
     primal_name = first(primal_invoke.args)
     pullback_expr = @strip_linenos quote
-        _make_pullback_for_non_differentiable(Val{$(tup_len_expr)}())
+        NonDiffPullback{$(tup_len_expr)}()
     end
 
     @gensym kwargs

test/rule_definition_tools.jl

@@ -42,27 +42,22 @@ end
 
 @testset "rule_definition_tools.jl" begin
     @testset "@non_differentiable" begin
-        @testset "`_make_pullback_for_non_differentiable`" begin
-            f = ChainRulesCore._make_pullback_for_non_differentiable
-            @testset "throws on invalid input" begin
-                @test_throws Exception f(Val(0.0))
-                @test_throws Exception f(Val(-1))
-            end
+        @testset "`NonDiffPullback`" begin
+            NDP = ChainRulesCore.NonDiffPullback
             @testset "identical objects" begin
                 for i in 0:5
-                    v = Val(i)
-                    @test f(v) === f(v)
+                    @test NDP{i}() === NDP{i}()
                 end
             end
             @testset "correctness" begin
                 for i in 0:5
                     expected = ntuple((_ -> NoTangent()), i)
-                    @test f(Val(i))(:arbitrary) === expected
+                    @test NDP{i}()(:arbitrary) === expected
                 end
             end
             @testset "dispatch" begin
                 for i in 0:5
-                    pullback = f(Val(i))
+                    pullback = NDP{i}()
                     @test_throws MethodError pullback()
                     @test_throws MethodError pullback(1, 2)
                 end