`@non_differentiable` should use identical pullbacks when possible

[nsajko]

`@non_differentiable` should use identical pullbacks when possible

[devmotion]

AFAICT the pullback should return a tuple of length N + 1. Alternatively with callable structs:

using ChainRulesCore

struct NonDiffPullback{N} end
(pb::NonDiffPullback{N})(::Vararg{Any,N}) where {N} = ntuple(Returns(NoTangent()), Val(N + 1))

[devmotion]

We can use Returns since the package depends on Compat >= 3.40 and it seems in version 3.35 support for Returns was added.

[nsajko]

I'm not good with macros so I probably won't be tackling this.

Actually, I think I got this. I think the issue can be fixed without relying on macros too much using make_pullback_for_non_differentiable:

using ChainRulesCore

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(::Vararg{Any,N})
        f = _ -> NoTangent()
        ntuple(f, Val(N))
    end
end

using Test

@testset "`make_pullback_for_non_differentiable`" begin
    f = 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 "identical objects" begin
        for i ∈ 0:5
            v = Val(i)
            @test f(v) === f(v)
        end
    end
    @testset "dispatch" begin
        pullback = f(Val(2))
        @test_throws MethodError pullback()
        @test_throws MethodError pullback(1)
        @test (NoTangent(), NoTangent()) === pullback(1, 2)
        @test_throws MethodError pullback(1, 2, 3)
    end
end

[nsajko]

Thanks, you're right, there's an off-by-one error. But note we can't use Returns until support for Julia v1.6 is dropped.