Fix CI on Julia 1.13 and nightly

.github/workflows/CI.yml

@@ -1,7 +1,7 @@
 name: CI
 on:
   push:
-    branches: [master]
+    branches: [main]
     tags: [v*]
   pull_request:
 

.github/workflows/IntegrationTest.yml

@@ -1,7 +1,7 @@
 name: IntegrationTest
 on:
   push:
-    branches: [master]
+    branches: [main]
     tags: [v*]
   pull_request:
 

.github/workflows/JuliaNightly.yml

@@ -2,7 +2,7 @@
 name: JuliaNightly
 on:
   push:
-    branches: [master]
+    branches: [main]
     tags: [v*]
   pull_request:
   schedule:

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "1.73.0"
+version = "1.73.1"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -27,8 +27,8 @@ Distributed = "1"
 FiniteDifferences = "0.12.20"
 GPUArraysCore = "0.1.0, 0.2"
 IrrationalConstants = "0.1.1, 0.2"
-JLArrays = "0.1"
-JuliaInterpreter = "0.8,0.9"
+JLArrays = "0.1, 0.4"
+JuliaInterpreter = "0.9, 0.12"
 LinearAlgebra = "1"
 Random = "1"
 RealDot = "0.1"

src/rulesets/Base/nondiff.jl

@@ -513,3 +513,6 @@ end
 @non_differentiable Threads.nthreads()
 @non_differentiable Threads.threadid()
 @non_differentiable Threads.threadid(::Task)
+if isdefined(Threads, :_spawn_set_thrpool)
+    @non_differentiable Threads._spawn_set_thrpool(::Task, ::Symbol)
+end

src/rulesets/Base/sort.jl

@@ -54,19 +54,22 @@ end
 
 function frule((_, ẋ), ::typeof(sortslices), x::AbstractArray; dims::Integer, kw...)
     p = sortperm(collect(eachslice(x; dims=dims)); kw...)
-    inds = ntuple(d -> d == dims ? p : (:), ndims(x))
+    inds = _sortslices_inds(x, dims, p)
     return x[inds...], ẋ[inds...]
 end
 
 function rrule(::typeof(sortslices), x::AbstractArray; dims::Integer, kw...)
     p = sortperm(collect(eachslice(x; dims=dims)); kw...)
-    inds = ntuple(d -> d == dims ? p : (:), ndims(x))
+    # the pullback does not capture the indices, whose type depends on the value of `dims`
     function sortslices_pullback(dy)
-        return (NoTangent(), ∇getindex(x, unthunk(dy), inds...))
+        return (NoTangent(), ∇getindex(x, unthunk(dy), _sortslices_inds(x, dims, p)...))
     end
-    return x[inds...], sortslices_pullback
+    return x[_sortslices_inds(x, dims, p)...], sortslices_pullback
 end
 
+# indices of `x` that select the slices along `dims` in the order `p`
+_sortslices_inds(x, dims, p) = ntuple(d -> d == dims ? p : (:), ndims(x))
+
 #####
 ##### `unique`
 #####

src/rulesets/LinearAlgebra/factorization.jl

@@ -20,6 +20,15 @@ const LU_NoPivot = VERSION >= v"1.7.0-DEV.1188" ? NoPivot : Val{false}
 
 const CHOLESKY_NoPivot = VERSION >= v"1.8.0-rc1" ? Union{NoPivot, Val{false}} : Val{false}
 
+# For fixed pivots, the factors are differentiable iff all pivots are nonzero, except
+# possibly the last pivot of square or wide A. `F.info` is the index of the first zero
+# pivot, negated for unpivoted factorizations since JuliaLang/julia#52957
+function _lu_isdifferentiable(F::LU, pivot)
+    m, n = size(F.factors)
+    info = pivot isa LU_NoPivot && VERSION >= v"1.11.0-DEV.1535" ? -F.info : F.info
+    return iszero(info) || (m <= n && info == m)
+end
+
 function frule(
     (_, Ȧ), ::typeof(lu!), A::StridedMatrix, pivot::Union{LU_RowMaximum,LU_NoPivot}; kwargs...
 )
@@ -28,34 +37,26 @@ function frule(
     ∂factors = pivot isa LU_RowMaximum ? ΔA[F.p, :] : ΔA
     m, n = size(∂factors)
     q = min(m, n)
-    if m == n  # square A
-        # minimal allocation computation of
-        # ∂L = L * tril(L \ (P * ΔA) / U, -1)
-        # ∂U = triu(L \ (P * ΔA) / U) * U
-        # ∂factors = ∂L + ∂U
-        L = UnitLowerTriangular(F.factors)
-        U = UpperTriangular(F.factors)
-        rdiv!(∂factors, U)
-        ldiv!(L, ∂factors)
-        ∂L = lmul!(L, tril(∂factors, -1))
-        ∂U = rmul!(triu(∂factors), U)
-        ∂factors .= ∂L .+ ∂U
-    elseif m < n  # wide A, system is [P*A1 P*A2] = [L*U1 L*U2]
-        L = UnitLowerTriangular(F.L)
-        U = F.U
+    if !_lu_isdifferentiable(F, pivot)
+        fill!(∂factors, NaN)
+    elseif m <= n  # square or wide A, system is [P*A1 P*A2] = [L*U1 L*U2]
+        # with B = L \ (P * ΔA) = [B1 B2]: ∂L = L * S and ∂U = triu(B - S * U) with
+        # S = tril(B1 / U1, -1), which only depends on U1[1:q-1, 1:q-1] (U[q, q] may be 0)
+        L = UnitLowerTriangular(view(F.factors, :, 1:q))
+        U1 = UpperTriangular(view(F.factors, :, 1:q))
         ldiv!(L, ∂factors)
         @views begin
             ∂factors1 = ∂factors[:, 1:q]
             ∂factors2 = ∂factors[:, (q + 1):end]
-            U1 = UpperTriangular(U[:, 1:q])
-            U2 = U[:, (q + 1):end]
+            U2 = F.factors[:, (q + 1):end]
         end
-        rdiv!(∂factors1, U1)
-        ∂L = tril(∂factors1, -1)
-        mul!(∂factors2, ∂L, U2, -1, 1)
-        lmul!(L, ∂L)
-        rmul!(triu!(∂factors1), U1)
-        ∂factors1 .+= ∂L
+        S = copy(∂factors1)
+        @views rdiv!(S[:, 1:(q - 1)], UpperTriangular(F.factors[1:(q - 1), 1:(q - 1)]))
+        tril!(S, -1)
+        mul!(∂factors2, S, U2, -1, 1)
+        ∂factors1 .-= rmul!(copy(S), U1)
+        triu!(∂factors1)
+        ∂factors1 .+= lmul!(L, S)
     else  # tall A, system is [P1*A; P2*A] = [L1*U; L2*U]
         L = F.L
         U = UpperTriangular(F.U)
@@ -86,30 +87,28 @@ function _lu_pullback(ΔF::Tangent, m, n, eltypeA, pivot, F)
     ∂factors = eltypeA <: Real ? real(Δfactors) : Δfactors
     ∂A = similar(factors)
     q = min(m, n)
-    if m == n  # square A
-        # ∂A = P' * (L' \ (tril(L' * ∂L, -1) + triu(∂U * U')) / U')
-        L = UnitLowerTriangular(factors)
-        U = UpperTriangular(factors)
-        ∂U = UpperTriangular(∂factors)
-        tril!(copyto!(∂A, ∂factors), -1)
-        lmul!(L', ∂A)
-        copyto!(UpperTriangular(∂A), UpperTriangular(∂U * U'))
-        rdiv!(∂A, U')
-        ldiv!(L', ∂A)
-    elseif m < n  # wide A, system is [P*A1 P*A2] = [L*U1 L*U2]
+    if !_lu_isdifferentiable(F, pivot)
+        fill!(∂A, NaN)
+    elseif m <= n  # square or wide A, system is [P*A1 P*A2] = [L*U1 L*U2]
+        # adjoint of the frule: with T = triu(∂factors) and
+        # ∂S = tril(L' * ∂factors1 - T * U', -1),
+        # ∂A = P' * (L' \ (T + [∂S[:, 1:q-1] / U[1:q-1, 1:q-1]' 0]))
         triu!(copyto!(∂A, ∂factors))
         @views begin
-            factors1 = factors[:, 1:q]
-            U2 = factors[:, (q + 1):end]
+            ∂factors1 = ∂factors[:, 1:q]
             ∂A1 = ∂A[:, 1:q]
             ∂A2 = ∂A[:, (q + 1):end]
-            ∂L = tril(∂factors[:, 1:q], -1)
+            L = UnitLowerTriangular(factors[:, 1:q])
+            U1 = UpperTriangular(factors[:, 1:q])
+            U2 = factors[:, (q + 1):end]
         end
-        L = UnitLowerTriangular(factors1)
-        U1 = UpperTriangular(factors1)
-        triu!(rmul!(∂A1, U1'))
-        ∂A1 .+= tril!(mul!(lmul!(L', ∂L), ∂A2, U2', -1, 1), -1)
-        rdiv!(∂A1, U1')
+        ∂S = lmul!(L', copy(∂factors1))
+        ∂S .-= rmul!(∂A1, U1')  # uses ∂A1 as scratch
+        mul!(∂S, ∂A2, U2', -1, 1)
+        tril!(∂S, -1)
+        @views rdiv!(∂S[:, 1:(q - 1)], UpperTriangular(factors[1:(q - 1), 1:(q - 1)])')
+        triu!(copyto!(∂A1, ∂factors1))
+        ∂A1 .+= ∂S
         ldiv!(L', ∂A)
     else  # tall A, system is [P1*A; P2*A] = [L1*U; L2*U]
         tril!(copyto!(∂A, ∂factors), -1)

src/rulesets/LinearAlgebra/norm.jl

@@ -257,11 +257,21 @@ end
 ##### `normalize`
 #####
 
+# `normalize!(y, p)` with precomputed `nrm == norm(y, p)`
+function _normalize!(y, nrm, p::Real=2)
+    if !issubnormal(nrm)
+        rmul!(y, inv(nrm))
+    else
+        normalize!(rmul!(y, inv(floatmin(nrm))), p)
+    end
+    return y
+end
+
 function rrule(::typeof(normalize), x::AbstractArray{<:Number}, p::Real)
     nrm, inner_pullback = rrule(norm, x, p)
     Ty = typeof(first(x) / nrm)
     y = copyto!(similar(x, Ty), x)
-    LinearAlgebra.__normalize!(y, nrm)
+    _normalize!(y, nrm, p)
     function normalize_pullback(Δy)
         invnrm = pinv(nrm)
         ∂nrm = -dot(y, Δy) * invnrm
@@ -277,7 +287,7 @@ function rrule(::typeof(normalize), x::AbstractArray{<:Number})
     nrm = LinearAlgebra.norm2(x)
     Ty = typeof(first(x) / nrm)
     y = copyto!(similar(x, Ty), x)
-    LinearAlgebra.__normalize!(y, nrm)
+    _normalize!(y, nrm)
     function normalize_pullback(ȳ)
         Δy = unthunk(ȳ)
         ∂x = (Δy .- realdot(y, Δy) .* y) .* pinv(nrm)

src/rulesets/LinearAlgebra/symmetric.jl

@@ -379,9 +379,14 @@ function _matfun(f, A::LinearAlgebra.RealHermSymComplexHerm)
     end
     fλ = first.(fλ_df_dλ)
     df_dλ = last.(unthunk.(fλ_df_dλ))
-    fA = (U * Diagonal(fλ)) * U'
+    fA = _matfun_primal(f, λ, U, fλ)
     Y = if eltype(A) <: Real
-        Symmetric(fA)
+        # match `f(A)`, which returns `Hermitian` since JuliaLang/julia#55951
+        @static if VERSION >= v"1.12.0-DEV.1322"
+            eltype(fλ) <: Complex ? Symmetric(fA) : _symhermtype(A)(fA)
+        else
+            Symmetric(fA)
+        end
     elseif eltype(fλ) <: Complex
         fA
     else
@@ -391,6 +396,23 @@ function _matfun(f, A::LinearAlgebra.RealHermSymComplexHerm)
     return Y, intermediates
 end
 
+# compute `f(A)` from its eigendecomposition exactly as `f(::RealHermSymComplexHerm)` does
+_matfun_primal(f, λ, U, fλ) = (U * Diagonal(fλ)) * U'
+# `exp` and positive semi-definite `sqrt` use `(U * √D) * (U * √D)'`
+# since JuliaLang/LinearAlgebra.jl#1359
+@static if VERSION >= v"1.13.0-DEV.887"
+    _matfun_primal(::typeof(exp), λ, U, fλ) = _psd_spectral_product(exp.(λ ./ 2), U)
+    function _matfun_primal(::typeof(sqrt), λ, U, fλ)
+        eltype(fλ) <: Real || return (U * Diagonal(fλ)) * U'
+        return _psd_spectral_product(fourthroot.(λ), U)
+    end
+    # local copy of `LinearAlgebra._psd_spectral_product` (internal)
+    function _psd_spectral_product(v, U)
+        Uv = U * Diagonal(v)
+        return Uv * Uv'
+    end
+end
+
 # Computes ∂Y = U * (P .* (U' * ΔA * U)) * U' with fewer allocations
 function _matfun_frechet(f, ΔA, A::LinearAlgebra.RealHermSymComplexHerm, Y, (λ, U, fλ, df_dλ))
     # We will overwrite tmp matrix several times to hold different values

test/rulesets/Base/array.jl

@@ -276,8 +276,15 @@ end
         @test_skip test_rrule(reverse, Tuple(rand(10)))  # Ambiguity in isapprox, https://github.com/JuliaDiff/ChainRulesTestUtils.jl/issues/229
     end
     @testset "Array" begin
+        # JLArrays ≥ 0.4 implements `reverse` without scalar indexing
+        jl_has_reverse = parentmodule(reverse, Tuple{JLArray{Float32,1}}) !== Base
+
         # Forward
-        @gpu_broken test_frule(reverse, rand(5))
+        if jl_has_reverse
+            @gpu test_frule(reverse, rand(5))
+        else
+            @gpu_broken test_frule(reverse, rand(5))
+        end
         test_frule(reverse, rand(5), 2, 4)
         test_frule(reverse, rand(5), fkwargs=(dims=1,))
         test_frule(reverse, rand(3,4), fkwargs=(dims=2,))
@@ -289,7 +296,11 @@ end
         test_frule(reverse!, rand(3,4), fkwargs=(dims=2,))
 
         # Reverse
-        @gpu_broken test_rrule(reverse, rand(5))
+        if jl_has_reverse
+            @gpu test_rrule(reverse, rand(5))
+        else
+            @gpu_broken test_rrule(reverse, rand(5))
+        end
         test_rrule(reverse, rand(5), 2, 4)
         test_rrule(reverse, rand(5), fkwargs=(dims=1,))
 

test/rulesets/Base/arraymath.jl

@@ -134,7 +134,8 @@
     end
 
     if VERSION > v"1.7.0-"
-        @eval using LinearAlgebra: mat_mat_scalar, mat_vec_scalar, StridedMaybeAdjOrTransMat
+        mat_mat_scalar = LinearAlgebra.mat_mat_scalar
+        mat_vec_scalar = LinearAlgebra.mat_vec_scalar
 
         @testset "3-arg *, $T" for T in [Float64, ComplexF64]
 

test/rulesets/Base/base.jl

@@ -224,6 +224,18 @@ end
         @test rrule(Base.depwarn, "message", :f) !== nothing
     end
 
+    if isdefined(Threads, :_spawn_set_thrpool)
+        @testset "Threads" begin
+            f = Threads._spawn_set_thrpool
+            _, ∂y = frule(
+                (NoTangent(), NoTangent(), NoTangent()), f, Task(() -> nothing), :default
+            )
+            @test ∂y isa NoTangent
+            _, back = rrule(f, Task(() -> nothing), :default)
+            @test back(1.0) == (NoTangent(), NoTangent(), NoTangent())
+        end
+    end
+
     @testset "literal_pow: $x^$p" for x in [-1.5, 0.0, 3.5], p in [2, 3]
         x == 0 && p < 0 && continue
         test_frule(Base.literal_pow, ^, x, Val(p))

test/rulesets/Base/broadcast.jl

@@ -110,8 +110,16 @@ BT1 = Broadcast.BroadcastStyle(Tuple)
 
             # These two test vararg rrule * rule:
             @gpu test_rrule(copy∘broadcasted, *, rand(3), rand(3), rand(3), rand(3), rand(3))
-            @gpu_broken test_rrule(copy∘broadcasted, *, rand(), rand(), rand(3), rand(3) .+ im, rand(4)')
-            # GPU error from dot(x::JLArray{Float32, 1}, y::JLArray{ComplexF32, 2})
+            if VERSION >= v"1.13"
+                @gpu test_rrule(
+                    copy∘broadcasted, *, rand(), rand(), rand(3), rand(3) .+ im, rand(4)'
+                )
+            else
+                @gpu_broken test_rrule(
+                    copy∘broadcasted, *, rand(), rand(), rand(3), rand(3) .+ im, rand(4)'
+                )
+                # GPU error from dot(x::JLArray{Float32, 1}, y::JLArray{ComplexF32, 2})
+            end
 
             @gpu test_rrule(copy∘broadcasted, Base.literal_pow, ^, rand(3), Val(2))
             @gpu test_rrule(copy∘broadcasted, Base.literal_pow, ^, rand(3) .+ im, Val(2))

test/rulesets/Base/indexing.jl

@@ -129,7 +129,7 @@ end
     end
 
     @testset "second derivatives: ∇getindex" begin
-        @eval using ChainRules: ∇getindex
+        ∇getindex = ChainRules.∇getindex
         # Forward, scalar result
         test_frule(∇getindex, rand(2, 3), rand(), 3)
         test_frule(∇getindex, rand(2, 3), rand()+im, 2, 1)