Merge remote-tracking branch 'origin/main' into gd/12

Author: gdalle

.github/workflows/Documentation.yml

@@ -35,7 +35,7 @@ jobs:
             dir: './DifferentiationInterfaceTest'
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
       - uses: julia-actions/setup-julia@v2
         with:
           version: '1'  # TODO: 1

.github/workflows/PreCommit.yml

@@ -18,7 +18,7 @@ jobs:
   pre-commit:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
       - uses: julia-actions/setup-julia@v2
       - uses: julia-actions/cache@v2
       - run: julia -e 'using Pkg; Pkg.add("JuliaFormatter")'

.github/workflows/Test.yml

@@ -67,7 +67,7 @@ jobs:
       JULIA_DI_TEST_GROUP: ${{ matrix.group }}
       JULIA_DI_PR_DRAFT: ${{ github.event.pull_request.draft }}
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
       - uses: julia-actions/setup-julia@v2
         with:
           version: ${{ matrix.version }}
@@ -126,7 +126,7 @@ jobs:
       JULIA_DIT_TEST_GROUP: ${{ matrix.group }}
       JULIA_DI_PR_DRAFT: ${{ github.event.pull_request.draft }}
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
       - uses: julia-actions/setup-julia@v2
         with:
           version: ${{ matrix.version }}

DifferentiationInterface/CHANGELOG.md

@@ -5,7 +5,22 @@ 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.5...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)
+
+### Fixed
+
+  - Put test deps into `test/Project.toml` ([#840](https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/840))
+  - Set up `pre-commit` ([#837](https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/837))
+
+### Fixed
+
+  - Put test deps into `test/Project.toml` ([#840](https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/840))
 
 ## [0.7.5](https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.4...DifferentiationInterface-v0.7.5)
 

DifferentiationInterface/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.7.5"
+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 -> basis(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 -> basis(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,11 @@ 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 = if x isa Number
+        oneunit(x)
+    else
+        basis(x)
+    end
     pushforward_prep = prepare_pushforward_nokwarg(
         strict, f, backend, x, (dx,), contexts...
     )
@@ -303,7 +307,11 @@ 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 = if x isa Number
+        oneunit(x)
+    else
+        basis(x)
+    end
     pushforward_prep = prepare_pushforward_nokwarg(
         strict, f!, y, backend, x, (dx,), contexts...
     )

DifferentiationInterface/src/first_order/pushforward.jl

@@ -290,7 +290,11 @@ 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 = if y isa Number
+        oneunit(y)
+    else
+        basis(y)
+    end
     pullback_prep = prepare_pullback_nokwarg(strict, f, backend, x, (dy,), contexts...)
     return PullbackPushforwardPrep(_sig, pullback_prep)
 end
@@ -306,7 +310,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 = basis(y)
     pullback_prep = prepare_pullback_nokwarg(strict, f!, y, backend, x, (dy,), contexts...)
     return PullbackPushforwardPrep(_sig, pullback_prep)
 end

DifferentiationInterface/src/second_order/hessian.jl

@@ -84,13 +84,15 @@ struct HVPGradientHessianPrep{
     BS<:BatchSizeSettings,
     S<:AbstractVector{<:NTuple},
     R<:AbstractVector{<:NTuple},
+    SE<:NTuple,
     E2<:HVPPrep,
     E1<:GradientPrep,
 } <: HessianPrep{SIG}
     _sig::Val{SIG}
     batch_size_settings::BS
     batched_seeds::S
     batched_results::R
+    seed_example::SE
     hvp_prep::E2
     gradient_prep::E1
 end
@@ -119,10 +121,17 @@ function _prepare_hessian_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]
-    hvp_prep = prepare_hvp_nokwarg(strict, f, backend, x, batched_seeds[1], contexts...)
+    seed_example = ntuple(b -> basis(x), Val(B))
+    hvp_prep = prepare_hvp_nokwarg(strict, f, backend, x, seed_example, contexts...)
     gradient_prep = prepare_gradient_nokwarg(strict, f, inner(backend), x, contexts...)
     return HVPGradientHessianPrep(
-        _sig, batch_size_settings, batched_seeds, batched_results, hvp_prep, gradient_prep
+        _sig,
+        batch_size_settings,
+        batched_seeds,
+        batched_results,
+        seed_example,
+        hvp_prep,
+        gradient_prep,
     )
 end
 
@@ -150,11 +159,11 @@ function hessian(
     contexts::Vararg{Context,C},
 ) where {F,SIG,B,aligned,C}
     check_prep(f, prep, backend, x, contexts...)
-    (; batch_size_settings, batched_seeds, hvp_prep) = prep
+    (; batch_size_settings, batched_seeds, seed_example, hvp_prep) = prep
     (; A, B_last) = batch_size_settings
 
     hvp_prep_same = prepare_hvp_same_point(
-        f, hvp_prep, backend, x, batched_seeds[1], contexts...
+        f, hvp_prep, backend, x, seed_example, contexts...
     )
 
     hess = mapreduce(hcat, eachindex(batched_seeds)) do a
@@ -178,11 +187,11 @@ function hessian!(
     contexts::Vararg{Context,C},
 ) where {F,SIG,B,C}
     check_prep(f, prep, backend, x, contexts...)
-    (; batch_size_settings, batched_seeds, batched_results, hvp_prep) = prep
+    (; batch_size_settings, batched_seeds, batched_results, seed_example, hvp_prep) = prep
     (; N) = batch_size_settings
 
     hvp_prep_same = prepare_hvp_same_point(
-        f, hvp_prep, backend, x, batched_seeds[1], contexts...
+        f, hvp_prep, backend, x, seed_example, contexts...
     )
 
     for a in eachindex(batched_seeds, batched_results)

DifferentiationInterface/src/utils/basis.jl

@@ -1,29 +1,39 @@
-"""
-    basis(a::AbstractArray, i)
+pre_basis(a::AbstractArray{T}) where {T} = fill!(similar(a), zero(T))
 
-Construct the `i`-th standard basis array in the vector space of `a`.
-"""
-function basis(a::AbstractArray{T}, i) where {T}
-    b = similar(a)
-    fill!(b, zero(T))
-    b[i] = oneunit(T)
+function post_basis(b::AbstractArray, a::AbstractArray)
     if ismutable_array(a)
         return b
     else
         return map(+, zero(a), b)
     end
 end
 
+"""
+    basis(a::AbstractArray, i)
+
+Construct the `i`-th standard basis array in the vector space of `a`.
+"""
+function basis(a::AbstractArray, i)
+    b = pre_basis(a)
+    b[i] = oneunit(eltype(b))
+    return post_basis(b, a)
+end
+
+# compatible with zero-length vectors
+function basis(a::AbstractArray)
+    b = pre_basis(a)
+    return post_basis(b, a)
+end
+
 """
     multibasis(a::AbstractArray, inds)
 
 Construct the sum of the `i`-th standard basis arrays in the vector space of `a` for all `i ∈ inds`.
 """
-function multibasis(a::AbstractArray{T}, inds) where {T}
-    b = similar(a)
-    fill!(b, zero(T))
+function multibasis(a::AbstractArray, inds)
+    b = pre_basis(a)
     for i in inds
-        b[i] = oneunit(T)
+        b[i] = oneunit(eltype(b))
     end
-    return ismutable_array(a) ? b : map(+, zero(a), b)
+    return post_basis(b, a)
 end