fix: improve support for empty inputs (still not guaranteed)

.github/workflows/Test.yml

@@ -25,7 +25,7 @@ jobs:
       actions: write
       contents: read
     strategy:
-      fail-fast: true  # TODO: toggle
+      fail-fast: false  # TODO: toggle
       matrix:
         version:
           - "1.10"

DifferentiationInterface/CHANGELOG.md

@@ -5,7 +5,11 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
-## [Unreleased](https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.6...main)
+## [Unreleased](https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.7...main)
+
+## [0.7.7](https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.6...DifferentiationInterface-v0.7.7)
+
+  - Improve support for empty inputs (still not guaranteed) ([#835](https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/835))
 
 ## [0.7.6](https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.5...DifferentiationInterface-v0.7.6)
 

DifferentiationInterface/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.7.6"
+version = "0.7.7"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

DifferentiationInterface/src/first_order/jacobian.jl

@@ -138,12 +138,14 @@ struct PushforwardJacobianPrep{
     BS<:BatchSizeSettings,
     S<:AbstractVector{<:NTuple},
     R<:AbstractVector{<:NTuple},
+    SE<:NTuple,
     E<:PushforwardPrep,
 } <: StandardJacobianPrep{SIG}
     _sig::Val{SIG}
     batch_size_settings::BS
     batched_seeds::S
     batched_results::R
+    seed_example::SE
     pushforward_prep::E
 end
 
@@ -152,12 +154,14 @@ struct PullbackJacobianPrep{
     BS<:BatchSizeSettings,
     S<:AbstractVector{<:NTuple},
     R<:AbstractVector{<:NTuple},
+    SE<:NTuple,
     E<:PullbackPrep,
 } <: StandardJacobianPrep{SIG}
     _sig::Val{SIG}
     batch_size_settings::BS
     batched_seeds::S
     batched_results::R
+    seed_example::SE
     pullback_prep::E
 end
 
@@ -211,11 +215,17 @@ function _prepare_jacobian_aux(
         ntuple(b -> seeds[1 + ((a - 1) * B + (b - 1)) % N], Val(B)) for a in 1:A
     ]
     batched_results = [ntuple(b -> similar(y), Val(B)) for _ in batched_seeds]
+    seed_example = ntuple(b -> zero(x), Val(B))
     pushforward_prep = prepare_pushforward_nokwarg(
-        strict, f_or_f!y..., backend, x, batched_seeds[1], contexts...
+        strict, f_or_f!y..., backend, x, seed_example, contexts...
     )
     return PushforwardJacobianPrep(
-        _sig, batch_size_settings, batched_seeds, batched_results, pushforward_prep
+        _sig,
+        batch_size_settings,
+        batched_seeds,
+        batched_results,
+        seed_example,
+        pushforward_prep,
     )
 end
 
@@ -236,11 +246,17 @@ function _prepare_jacobian_aux(
         ntuple(b -> seeds[1 + ((a - 1) * B + (b - 1)) % N], Val(B)) for a in 1:A
     ]
     batched_results = [ntuple(b -> similar(x), Val(B)) for _ in batched_seeds]
+    seed_example = ntuple(b -> zero(y), Val(B))
     pullback_prep = prepare_pullback_nokwarg(
-        strict, f_or_f!y..., backend, x, batched_seeds[1], contexts...
+        strict, f_or_f!y..., backend, x, seed_example, contexts...
     )
     return PullbackJacobianPrep(
-        _sig, batch_size_settings, batched_seeds, batched_results, pullback_prep
+        _sig,
+        batch_size_settings,
+        batched_seeds,
+        batched_results,
+        seed_example,
+        pullback_prep,
     )
 end
 
@@ -363,11 +379,11 @@ function _jacobian_aux(
     x,
     contexts::Vararg{Context,C},
 ) where {FY,SIG,B,aligned,C}
-    (; batch_size_settings, batched_seeds, pushforward_prep) = prep
+    (; batch_size_settings, batched_seeds, seed_example, pushforward_prep) = prep
     (; A, B_last) = batch_size_settings
 
     pushforward_prep_same = prepare_pushforward_same_point(
-        f_or_f!y..., pushforward_prep, backend, x, batched_seeds[1], contexts...
+        f_or_f!y..., pushforward_prep, backend, x, seed_example, contexts...
     )
 
     jac = mapreduce(hcat, eachindex(batched_seeds)) do a
@@ -419,11 +435,11 @@ function _jacobian_aux(
     x,
     contexts::Vararg{Context,C},
 ) where {FY,SIG,B,aligned,C}
-    (; batch_size_settings, batched_seeds, pullback_prep) = prep
+    (; batch_size_settings, batched_seeds, seed_example, pullback_prep) = prep
     (; A, B_last) = batch_size_settings
 
     pullback_prep_same = prepare_pullback_same_point(
-        f_or_f!y..., prep.pullback_prep, backend, x, batched_seeds[1], contexts...
+        f_or_f!y..., pullback_prep, backend, x, seed_example, contexts...
     )
 
     jac = mapreduce(vcat, eachindex(batched_seeds)) do a
@@ -451,11 +467,13 @@ function _jacobian_aux!(
     x,
     contexts::Vararg{Context,C},
 ) where {FY,SIG,B,C}
-    (; batch_size_settings, batched_seeds, batched_results, pushforward_prep) = prep
+    (;
+        batch_size_settings, batched_seeds, batched_results, seed_example, pushforward_prep
+    ) = prep
     (; N) = batch_size_settings
 
     pushforward_prep_same = prepare_pushforward_same_point(
-        f_or_f!y..., pushforward_prep, backend, x, batched_seeds[1], contexts...
+        f_or_f!y..., pushforward_prep, backend, x, seed_example, contexts...
     )
 
     for a in eachindex(batched_seeds, batched_results)
@@ -487,11 +505,12 @@ function _jacobian_aux!(
     x,
     contexts::Vararg{Context,C},
 ) where {FY,SIG,B,C}
-    (; batch_size_settings, batched_seeds, batched_results, pullback_prep) = prep
+    (; batch_size_settings, batched_seeds, batched_results, seed_example, pullback_prep) =
+        prep
     (; N) = batch_size_settings
 
     pullback_prep_same = prepare_pullback_same_point(
-        f_or_f!y..., pullback_prep, backend, x, batched_seeds[1], contexts...
+        f_or_f!y..., pullback_prep, backend, x, seed_example, contexts...
     )
 
     for a in eachindex(batched_seeds, batched_results)

DifferentiationInterface/src/first_order/pullback.jl

@@ -285,7 +285,7 @@ function _prepare_pullback_aux(
     contexts::Vararg{Context,C};
 ) where {F,C}
     _sig = signature(f, backend, x, ty, contexts...; strict)
-    dx = x isa Number ? oneunit(x) : basis(x, first(CartesianIndices(x)))
+    dx = zero(x)
     pushforward_prep = prepare_pushforward_nokwarg(
         strict, f, backend, x, (dx,), contexts...
     )
@@ -303,7 +303,7 @@ function _prepare_pullback_aux(
     contexts::Vararg{Context,C};
 ) where {F,C}
     _sig = signature(f!, y, backend, x, ty, contexts...; strict)
-    dx = x isa Number ? oneunit(x) : basis(x, first(CartesianIndices(x)))
+    dx = zero(x)
     pushforward_prep = prepare_pushforward_nokwarg(
         strict, f!, y, backend, x, (dx,), contexts...
     )

DifferentiationInterface/src/first_order/pushforward.jl

@@ -290,7 +290,7 @@ function _prepare_pushforward_aux(
 ) where {F,C}
     _sig = signature(f, backend, x, tx, contexts...; strict)
     y = f(x, map(unwrap, contexts)...)
-    dy = y isa Number ? oneunit(y) : basis(y, first(CartesianIndices(y)))
+    dy = zero(y)
     pullback_prep = prepare_pullback_nokwarg(strict, f, backend, x, (dy,), contexts...)
     return PullbackPushforwardPrep(_sig, pullback_prep)
 end
@@ -306,7 +306,7 @@ function _prepare_pushforward_aux(
     contexts::Vararg{Context,C};
 ) where {F,C}
     _sig = signature(f!, y, backend, x, tx, contexts...; strict)
-    dy = y isa Number ? oneunit(y) : basis(y, first(CartesianIndices(y)))
+    dy = zero(y)
     pullback_prep = prepare_pullback_nokwarg(strict, f!, y, backend, x, (dy,), contexts...)
     return PullbackPushforwardPrep(_sig, pullback_prep)
 end

DifferentiationInterface/src/utils/batchsize.jl

@@ -6,7 +6,7 @@ Configuration for the batch size deduced from a backend and a sample array of le
 # Type parameters
 
   - `B::Int`: batch size
-  - `singlebatch::Bool`: whether `B == N` (`B > N` is not allowed)
+  - `singlebatch::Bool`: whether `B == N` (`B > N` is only allowed when `N == 0`)
   - `aligned::Bool`: whether `N % B == 0`
 
 # Fields
@@ -22,22 +22,26 @@ struct BatchSizeSettings{B,singlebatch,aligned}
 end
 
 function BatchSizeSettings{B,singlebatch,aligned}(N::Integer) where {B,singlebatch,aligned}
-    B > N && throw(ArgumentError("Batch size $B larger than input size $N"))
-    A = div(N, B, RoundUp)
-    B_last = N % B
+    B > N > 0 && throw(ArgumentError("Batch size $B larger than input size $N"))
+    if B == N == 0
+        A = B_last = 0
+    else
+        A = div(N, B, RoundUp)
+        B_last = N % B
+    end
     return BatchSizeSettings{B,singlebatch,aligned}(N, A, B_last)
 end
 
 function BatchSizeSettings{B}(::Val{N}) where {B,N}
     singlebatch = B == N
-    aligned = N % B == 0
+    aligned = (B == N == 0) || (N % B == 0)
     return BatchSizeSettings{B,singlebatch,aligned}(N)
 end
 
 function BatchSizeSettings{B}(N::Integer) where {B}
     # type-unstable
     singlebatch = B == N
-    aligned = N % B == 0
+    aligned = (B == N == 0) || (N % B == 0)
     return BatchSizeSettings{B,singlebatch,aligned}(N)
 end
 
@@ -123,7 +127,9 @@ Reproduces the heuristic from ForwardDiff to minimize
 Source: https://github.com/JuliaDiff/ForwardDiff.jl/blob/ec74fbc32b10bbf60b3c527d8961666310733728/src/prelude.jl#L19-L29
 """
 function reasonable_batchsize(N::Integer, Bmax::Integer)
-    if N <= Bmax
+    if N == 0
+        return 1
+    elseif N <= Bmax
         return N
     else
         A = div(N, Bmax, RoundUp)

DifferentiationInterface/test/Back/Enzyme/test.jl

@@ -201,3 +201,11 @@ end
         @test occursin("DifferentiationInterface", msg)
     end
 end
+
+@testset "Empty arrays" begin
+    test_differentiation(
+        [AutoEnzyme(; mode=Enzyme.Forward), AutoEnzyme(; mode=Enzyme.Reverse)],
+        empty_scenarios();
+        excluded=[:jacobian],
+    )
+end;

DifferentiationInterface/test/Core/Internals/batchsize.jl

@@ -13,6 +13,7 @@ using Test
 BSS = BatchSizeSettings
 
 @testset "Default" begin
+    @test (@inferred pick_batchsize(AutoZygote(), zeros(0))) isa BSS{1,false,true}
     @test (@inferred pick_batchsize(AutoZygote(), zeros(2))) isa BSS{1,false,true}
     @test (@inferred pick_batchsize(AutoZygote(), zeros(100))) isa BSS{1,false,true}
     @test_throws ArgumentError pick_batchsize(AutoSparse(AutoZygote()), zeros(2))
@@ -25,11 +26,14 @@ BSS = BatchSizeSettings
 end
 
 @testset "SimpleFiniteDiff (adaptive)" begin
+    @test (pick_batchsize(AutoSimpleFiniteDiff(), zeros(0))) isa BSS{1,false,true}
     @test (pick_batchsize(AutoSimpleFiniteDiff(), zeros(2))) isa BSS{2,true,true}
     @test (pick_batchsize(AutoSimpleFiniteDiff(), zeros(6))) isa BSS{6,true,true}
     @test (pick_batchsize(AutoSimpleFiniteDiff(), zeros(12))) isa BSS{12,true,true}
     @test (pick_batchsize(AutoSimpleFiniteDiff(), zeros(24))) isa BSS{12,false,true}
     @test (pick_batchsize(AutoSimpleFiniteDiff(), zeros(100))) isa BSS{12,false,false}
+    @test (@inferred pick_batchsize(AutoSimpleFiniteDiff(), @SVector(zeros(0)))) isa
+        BSS{0,true,true}
     @test (@inferred pick_batchsize(AutoSimpleFiniteDiff(), @SVector(zeros(2)))) isa
         BSS{2,true,true}
     @test (@inferred pick_batchsize(AutoSimpleFiniteDiff(), @SVector(zeros(6)))) isa

DifferentiationInterface/test/Core/ZeroBackends/test.jl

@@ -1,6 +1,7 @@
 using DifferentiationInterface
 using DifferentiationInterface: AutoZeroForward, AutoZeroReverse
 using DifferentiationInterfaceTest
+using LinearAlgebra
 using ComponentArrays: ComponentArrays
 using JLArrays: JLArrays
 using SparseMatrixColorings
@@ -50,3 +51,9 @@ end
         logging=LOGGING,
     )
 end
+
+@testset "Empty arrays" begin
+    test_differentiation(
+        [AutoZeroForward(), AutoZeroReverse()], empty_scenarios(); excluded=[:jacobian]
+    )
+end;

DifferentiationInterface/test/testutils.jl

@@ -10,7 +10,8 @@ using DifferentiationInterfaceTest:
     complex_sparse_scenarios,
     static_scenarios,
     component_scenarios,
-    gpu_scenarios
+    gpu_scenarios,
+    empty_scenarios
 
 function MyAutoSparse(backend::AbstractADType)
     return AutoSparse(

DifferentiationInterfaceTest/CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+  - Improve support for empty inputs (still not guaranteed) ([#835](https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/835))
   - Compute Scenario results with a reference backend ([#839](https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/839))
 
 ### Fixed

DifferentiationInterfaceTest/src/DifferentiationInterfaceTest.jl

@@ -93,7 +93,7 @@ using DifferentiationInterface: Rewrap, Context, Constant, Cache, ConstantOrCach
 using DifferentiationInterface: PreparationMismatchError
 using DocStringExtensions: TYPEDFIELDS, TYPEDSIGNATURES
 using JET: @test_opt
-using LinearAlgebra: Adjoint, Diagonal, Transpose, dot, parent
+using LinearAlgebra: Adjoint, Diagonal, Transpose, I, dot, parent
 using ProgressMeter: ProgressUnknown, next!
 using Random: AbstractRNG, default_rng, rand!
 using SparseArrays:
@@ -124,6 +124,7 @@ include("scenarios/default.jl")
 include("scenarios/sparse.jl")
 include("scenarios/complex.jl")
 include("scenarios/allocfree.jl")
+include("scenarios/empty.jl")
 include("scenarios/extensions.jl")
 
 include("tests/correctness_eval.jl")

DifferentiationInterfaceTest/src/scenarios/empty.jl

@@ -0,0 +1,16 @@
+make_empty(t::Number) = typeof(t)[]
+function make_empty!(y::AbstractArray, t::Number)
+    @assert isempty(y)
+    return nothing
+end
+
+function empty_scenarios()
+    scens = Scenario[
+        Scenario{:derivative,:out}(make_empty, 1.0; res1=Float64[]),
+        Scenario{:derivative,:out}(make_empty!, Float64[], 1.0; res1=Float64[]),
+        Scenario{:gradient,:out}(sum, Float64[]; res1=Float64[]),
+        Scenario{:jacobian,:out}(copy, Float64[]; res1=float.(I(0))),
+        Scenario{:jacobian,:out}(copyto!, Float64[], Float64[]; res1=float.(I(0))),
+    ]
+    return scens
+end

DifferentiationInterfaceTest/test/runtests.jl

@@ -12,7 +12,8 @@ using DifferentiationInterfaceTest:
     complex_sparse_scenarios,
     static_scenarios,
     component_scenarios,
-    gpu_scenarios
+    gpu_scenarios,
+    empty_scenarios
 
 GROUP = get(ENV, "JULIA_DIT_TEST_GROUP", "All")
 

DifferentiationInterfaceTest/test/standard.jl

@@ -28,6 +28,13 @@ test_differentiation(
     logging=LOGGING,
 )
 
+test_differentiation(
+    [AutoForwardDiff()], empty_scenarios(); excluded=[:gradient], logging=LOGGING
+)
+test_differentiation(
+    [AutoFiniteDiff()], empty_scenarios(); excluded=[:jacobian], logging=LOGGING
+)
+
 ## Sparse
 
 sparse_backend = AutoSparse(

DifferentiationInterfaceTest/test/zero_backends.jl

@@ -86,7 +86,9 @@ end
             logging=LOGGING,
         ),
     )
-    @test all(iszero, data_allocfree[!, :allocs])
+    @testset "$(collect(row[1:4]))" for row in collect(eachrow(data_allocfree))
+        @test row[:allocs] == 0
+    end
 end
 
 test_differentiation(