From ee828d03ad347b85a1b7493e4e42d122ce8b5123 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Wed, 14 Oct 2026 10:38:45 +0200 Subject: [PATCH 01/12] Fix CI on Julia 1.13 and nightly - `@gpu` test helper: splice the name `_gpu_test` instead of the function value. Since Julia 1.13.0-DEV.300 (JuliaLang/julia#57839) `@test` lowers to `eval_test_function(..., quoted_func::Union{Expr,Symbol}, ...)`, so every `@gpu` test errored with a `MethodError`. - `_matfun`: real-valued functions of `Hermitian{<:Real}` return `Hermitian` since Julia 1.12.0-DEV.1322 (JuliaLang/julia#55951); match the primal type. - Mark `Threads._spawn_set_thrpool` non-differentiable (hit by `Threads.@spawn` in Zygote on Julia 1.13). - Test-only compat: `JuliaInterpreter = "0.9, 0.12"`; 0.9 no longer precompiles on nightly. - Workflows: run on pushes to `main` (was `master`, which doesn't exist). Co-Authored-By: Claude Opus 5.5 --- .github/workflows/CI.yml | 2 +- .github/workflows/IntegrationTest.yml | 2 +- .github/workflows/JuliaNightly.yml | 2 +- Project.toml | 4 ++-- src/rulesets/Base/nondiff.jl | 3 +++ src/rulesets/LinearAlgebra/symmetric.jl | 7 ++++++- test/rulesets/LinearAlgebra/symmetric.jl | 7 +++++-- test/test_helpers.jl | 5 +++-- 8 files changed, 22 insertions(+), 10 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 1ab6f2af2..16b076a17 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -1,7 +1,7 @@ name: CI on: push: - branches: [master] + branches: [main] tags: [v*] pull_request: diff --git a/.github/workflows/IntegrationTest.yml b/.github/workflows/IntegrationTest.yml index e8b24668b..c10a5ab71 100644 --- a/.github/workflows/IntegrationTest.yml +++ b/.github/workflows/IntegrationTest.yml @@ -1,7 +1,7 @@ name: IntegrationTest on: push: - branches: [master] + branches: [main] tags: [v*] pull_request: diff --git a/.github/workflows/JuliaNightly.yml b/.github/workflows/JuliaNightly.yml index 260c52fab..b425ab3b9 100644 --- a/.github/workflows/JuliaNightly.yml +++ b/.github/workflows/JuliaNightly.yml @@ -2,7 +2,7 @@ name: JuliaNightly on: push: - branches: [master] + branches: [main] tags: [v*] pull_request: schedule: diff --git a/Project.toml b/Project.toml index ff52dd8bb..c7172cc90 100644 --- a/Project.toml +++ b/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" @@ -28,7 +28,7 @@ FiniteDifferences = "0.12.20" GPUArraysCore = "0.1.0, 0.2" IrrationalConstants = "0.1.1, 0.2" JLArrays = "0.1" -JuliaInterpreter = "0.8,0.9" +JuliaInterpreter = "0.9, 0.12" LinearAlgebra = "1" Random = "1" RealDot = "0.1" diff --git a/src/rulesets/Base/nondiff.jl b/src/rulesets/Base/nondiff.jl index 3da755aad..184413a5e 100644 --- a/src/rulesets/Base/nondiff.jl +++ b/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 diff --git a/src/rulesets/LinearAlgebra/symmetric.jl b/src/rulesets/LinearAlgebra/symmetric.jl index cdadca6c3..bd6ea60ad 100644 --- a/src/rulesets/LinearAlgebra/symmetric.jl +++ b/src/rulesets/LinearAlgebra/symmetric.jl @@ -381,7 +381,12 @@ function _matfun(f, A::LinearAlgebra.RealHermSymComplexHerm) df_dλ = last.(unthunk.(fλ_df_dλ)) fA = (U * Diagonal(fλ)) * U' 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 diff --git a/test/rulesets/LinearAlgebra/symmetric.jl b/test/rulesets/LinearAlgebra/symmetric.jl index 593b82148..d824556b6 100644 --- a/test/rulesets/LinearAlgebra/symmetric.jl +++ b/test/rulesets/LinearAlgebra/symmetric.jl @@ -314,6 +314,9 @@ end end + # wrapper of real-valued `f(A)` for real `A` (`TA` since JuliaLang/julia#55951) + realfun_wrapper(TA) = VERSION >= v"1.12.0-DEV.1322" ? TA : Symmetric + @testset "$(f)(::$TA{<:$T})" for f in (exp, log, sqrt, cos, sin, tan, cosh, sinh, tanh, acos, asin, atan, acosh, asinh, atanh), TA in (Symmetric, Hermitian), @@ -329,7 +332,7 @@ Y_ad, ∂Y_ad = @maybe_inferred frule((ZeroTangent(), ΔA), f, A) else TY = T∂Y = if T <: Real - Union{Symmetric{Complex{T}},Symmetric{T}} + Union{Symmetric{Complex{T}},realfun_wrapper(TA){T}} else Union{Matrix{T},Hermitian{T}} end @@ -382,7 +385,7 @@ Y_ad, back = @maybe_inferred rrule(f, A) else TY = if T <: Real - Union{Symmetric{Complex{T}},Symmetric{T}} + Union{Symmetric{Complex{T}},realfun_wrapper(TA){T}} else Union{Matrix{T},Hermitian{T}} end diff --git a/test/test_helpers.jl b/test/test_helpers.jl index 1aa85d52c..2aebaccdc 100644 --- a/test/test_helpers.jl +++ b/test/test_helpers.jl @@ -42,10 +42,11 @@ end function _gpu_macro(ex, broken::Bool, __source__) Meta.isexpr(ex, :call) && ex.args[1] in (:test_rrule, :test_frule, :rrule, :frule) || error("@gpu doesn't understand this input") + # `@test` requires a Symbol, not a function (Julia ≥ 1.13.0-DEV.300) ex2 = if Meta.isexpr(ex.args[2], :parameters) - :($_gpu_test($(ex.args[2]), $(ex.args[1]), $(ex.args[3:end]...))) + :(_gpu_test($(ex.args[2]), $(ex.args[1]), $(ex.args[3:end]...))) else - :($_gpu_test($(ex.args[1]), $(ex.args[2:end]...))) + :(_gpu_test($(ex.args[1]), $(ex.args[2:end]...))) end return if broken Expr(:block, ex, Expr(:macrocall, Symbol("@test_broken"), __source__, ex2)) |> esc From 5d88aea1601f03242ac2902d60192bab529630a6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Wed, 14 Oct 2026 14:46:25 +0200 Subject: [PATCH 02/12] Fix remaining CI failures on Julia 1.13 and nightly - Test-only compat: `JLArrays = "0.1, 0.4"`. With `@gpu` tests running again, GPUArrays 10 (JLArrays 0.1) hits scalar indexing in Julia >= 1.12 LinearAlgebra (`tr(::StridedMatrix)`, 3x3 `matmul3x3!`, Diagonal `mul!`); GPUArrays 12.1 (JLArrays 0.4) overrides these. - Tests: avoid `@eval using`/`@eval struct` inside interpreted testsets. JuliaInterpreter 0.12 respects world age, so those bindings were invisible. - `sortslices` rrule: don't capture the index tuple, whose type depends on the value of `dims`; the rrule is now inferrable. - broadcast.jl: vararg `*` broadcast works on GPU since Julia 1.13 (still broken on 1.6, 1.10, 1.11, 1.12). - `normalize` rrules: replace `LinearAlgebra.__normalize!`, removed in JuliaLang/LinearAlgebra.jl#1634, by a local copy of the upstream logic. - structured.jl: skip `Adjoint(::Adjoint)`/`Transpose(::Transpose)` where they throw (JuliaLang/LinearAlgebra.jl#1739). Co-Authored-By: Claude Opus 5.5 --- Project.toml | 2 +- src/rulesets/Base/sort.jl | 4 ++-- src/rulesets/LinearAlgebra/norm.jl | 14 ++++++++++++-- test/rulesets/Base/arraymath.jl | 3 ++- test/rulesets/Base/broadcast.jl | 8 ++++++-- test/rulesets/Base/indexing.jl | 2 +- test/rulesets/LinearAlgebra/structured.jl | 23 +++++++++++++++++------ test/rulesets/Statistics/statistics.jl | 5 +++-- 8 files changed, 44 insertions(+), 17 deletions(-) diff --git a/Project.toml b/Project.toml index c7172cc90..d4e54cad2 100644 --- a/Project.toml +++ b/Project.toml @@ -27,7 +27,7 @@ Distributed = "1" FiniteDifferences = "0.12.20" GPUArraysCore = "0.1.0, 0.2" IrrationalConstants = "0.1.1, 0.2" -JLArrays = "0.1" +JLArrays = "0.1, 0.4" JuliaInterpreter = "0.9, 0.12" LinearAlgebra = "1" Random = "1" diff --git a/src/rulesets/Base/sort.jl b/src/rulesets/Base/sort.jl index 0805da91f..de2bd4ece 100644 --- a/src/rulesets/Base/sort.jl +++ b/src/rulesets/Base/sort.jl @@ -60,11 +60,11 @@ 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)) function sortslices_pullback(dy) + inds = ntuple(d -> d == dims ? p : (:), ndims(x)) return (NoTangent(), ∇getindex(x, unthunk(dy), inds...)) end - return x[inds...], sortslices_pullback + return x[ntuple(d -> d == dims ? p : (:), ndims(x))...], sortslices_pullback end ##### diff --git a/src/rulesets/LinearAlgebra/norm.jl b/src/rulesets/LinearAlgebra/norm.jl index 89d280134..f3973e0a1 100644 --- a/src/rulesets/LinearAlgebra/norm.jl +++ b/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) diff --git a/test/rulesets/Base/arraymath.jl b/test/rulesets/Base/arraymath.jl index 2099b9b3b..8b8104f2b 100644 --- a/test/rulesets/Base/arraymath.jl +++ b/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] diff --git a/test/rulesets/Base/broadcast.jl b/test/rulesets/Base/broadcast.jl index de7516bfb..9db97f048 100644 --- a/test/rulesets/Base/broadcast.jl +++ b/test/rulesets/Base/broadcast.jl @@ -110,8 +110,12 @@ 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)) diff --git a/test/rulesets/Base/indexing.jl b/test/rulesets/Base/indexing.jl index d947567e5..e0e17237c 100644 --- a/test/rulesets/Base/indexing.jl +++ b/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) diff --git a/test/rulesets/LinearAlgebra/structured.jl b/test/rulesets/LinearAlgebra/structured.jl index c46460f46..db7148f06 100644 --- a/test/rulesets/LinearAlgebra/structured.jl +++ b/test/rulesets/LinearAlgebra/structured.jl @@ -1,3 +1,10 @@ +# `Adjoint(::Adjoint)` and `Transpose(::Transpose)` throw since JuliaLang/LinearAlgebra.jl#1739 +const DOUBLE_WRAP_ALLOWED = try + Adjoint(Adjoint([1.0])) isa Adjoint +catch + false +end + @testset "Structured Matrices" begin @testset "/ and \\ on Square Matrixes" begin @testset "//, $T on the RHS" for T in (Diagonal, UpperTriangular, LowerTriangular) @@ -122,14 +129,18 @@ @test pb(ȳ_mat) == pb(ȳ_composite) end - @testset "$f(::Adjoint{$T, Vector{$T})" begin - a = randn(T, n)' - test_rrule(f, a; output_tangent=randn(T, n)) + if f !== Adjoint || DOUBLE_WRAP_ALLOWED + @testset "$f(::Adjoint{$T, Vector{$T})" begin + a = randn(T, n)' + test_rrule(f, a; output_tangent=randn(T, n)) + end end - @testset "$f(::Transpose{$T, Vector{$T})" begin - a = transpose(randn(T, n)) - test_rrule(f, a; output_tangent=randn(T, n)) + if f !== Transpose || DOUBLE_WRAP_ALLOWED + @testset "$f(::Transpose{$T, Vector{$T})" begin + a = transpose(randn(T, n)) + test_rrule(f, a; output_tangent=randn(T, n)) + end end end @testset "$T" for T in (UpperTriangular, LowerTriangular) diff --git a/test/rulesets/Statistics/statistics.jl b/test/rulesets/Statistics/statistics.jl index 1b41fac0b..657542d69 100644 --- a/test/rulesets/Statistics/statistics.jl +++ b/test/rulesets/Statistics/statistics.jl @@ -1,3 +1,6 @@ +struct DummyWeights <: AbstractVector{Float64} # DummyType that looks like StatsBase's Weights types +end + @testset "mean" begin @testset "mean(x)" begin @gpu test_rrule(mean, randn(9)) @@ -30,8 +33,6 @@ end @testset "Regression Test against StatsBase-like Weighted Mean" begin - @eval struct DummyWeights <: AbstractVector{Float64} # DummyType that looks like StatsBase's Weights types - end # This should return nothing as we have no rule for this. (we opted opt) @test nothing == rrule(ChainRulesTestUtils.TestConfig(), mean, [1.0, 2.0], DummyWeights()) end From e59fd176582997cb1795f668f84148985118c545 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Wed, 14 Oct 2026 15:09:10 +0200 Subject: [PATCH 03/12] Apply batched suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- test/rulesets/Base/broadcast.jl | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/test/rulesets/Base/broadcast.jl b/test/rulesets/Base/broadcast.jl index 9db97f048..f486852a1 100644 --- a/test/rulesets/Base/broadcast.jl +++ b/test/rulesets/Base/broadcast.jl @@ -111,9 +111,13 @@ 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)) if VERSION >= v"1.13" - @gpu test_rrule(copy∘broadcasted, *, rand(), rand(), rand(3), rand(3) .+ im, rand(4)') + @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_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 From e39a6e9910e62150ae093454523d0776e5c53a3f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Wed, 14 Oct 2026 23:41:13 +0200 Subject: [PATCH 04/12] Match Base's PSD `exp`/`sqrt` and fix `reverse` GPU test with JLArrays 0.4 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Since JuliaLang/LinearAlgebra.jl#1359 (Julia 1.13.0-DEV.887), `exp` and positive semi-definite `sqrt` of `Hermitian`/`Symmetric` matrices are computed as `(U * √D) * (U * √D)'`, so `_matfun` no longer reproduced `f(A)` exactly. JLArrays 0.4 implements `reverse` without scalar indexing, so the `@gpu_broken` tests passed unexpectedly. Co-Authored-By: Claude Opus 5.5 --- src/rulesets/LinearAlgebra/symmetric.jl | 19 ++++++++++++++++++- test/rulesets/Base/array.jl | 15 +++++++++++++-- 2 files changed, 31 insertions(+), 3 deletions(-) diff --git a/src/rulesets/LinearAlgebra/symmetric.jl b/src/rulesets/LinearAlgebra/symmetric.jl index bd6ea60ad..23e4b2941 100644 --- a/src/rulesets/LinearAlgebra/symmetric.jl +++ b/src/rulesets/LinearAlgebra/symmetric.jl @@ -379,7 +379,7 @@ 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 # match `f(A)`, which returns `Hermitian` since JuliaLang/julia#55951 @static if VERSION >= v"1.12.0-DEV.1322" @@ -396,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 diff --git a/test/rulesets/Base/array.jl b/test/rulesets/Base/array.jl index 9a715d623..acb207b08 100644 --- a/test/rulesets/Base/array.jl +++ b/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,)) From a633c1e9fee7480b205dc0e7d831b15fe98ce69d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Thu, 15 Oct 2026 10:08:33 +0200 Subject: [PATCH 05/12] Run testsets compiled on Julia 1.14 JuliaInterpreter 0.12 fails on the new `foreigncall` form and looks up methods with `Type{T}` instead of `Core.TypeEgal{T}` signatures. Co-Authored-By: Claude Opus 5.5 --- test/runtests.jl | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/test/runtests.jl b/test/runtests.jl index 768f7c208..22791e3cd 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -24,13 +24,16 @@ using JuliaInterpreter union!(JuliaInterpreter.compiled_modules, Any[Base, Base.Broadcast, LinearAlgebra, Random, StaticArrays, Statistics]) +# JuliaInterpreter (0.12) does not support Julia 1.14 yet +const INTERPRET_TESTSETS = VERSION < v"1.14.0-DEV" + Random.seed!(1) # Set seed that all testsets should reset to. function include_test(path) if isempty(ARGS) || any(occursin(a, path) for a in ARGS) println("Testing $path:") # print so TravisCI doesn't timeout due to no output @time Base.include(@__MODULE__(), path) do ex - Meta.isexpr(ex, :macrocall) && ex.args[1] == Symbol("@testset") || return ex + INTERPRET_TESTSETS && Meta.isexpr(ex, :macrocall) && ex.args[1] == Symbol("@testset") || return ex return :(@interpret (() -> $ex)()) # interpret testsets using JuliaInterpreter end else From d67f02e5bbe76e22fb4868a1a2fe4f99405b263a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Thu, 15 Oct 2026 15:40:46 +0200 Subject: [PATCH 06/12] Fix `lu` rules for singular matrices MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With JuliaLang/LinearAlgebra.jl#1529, triangular solves throw for zero diagonal entries, so the `lu!` frule failed for singular A with `check=false` on Julia 1.14. For square and wide A the rules now only invert the leading (q-1)×(q-1) block of U, so they compute the derivative, which exists if only the last pivot is zero. Otherwise, the factors are not differentiable and the rules return NaN. Co-Authored-By: Claude Opus 5.5 --- src/rulesets/LinearAlgebra/factorization.jl | 83 ++++++++++---------- test/rulesets/LinearAlgebra/factorization.jl | 53 +++++++++++-- 2 files changed, 89 insertions(+), 47 deletions(-) diff --git a/src/rulesets/LinearAlgebra/factorization.jl b/src/rulesets/LinearAlgebra/factorization.jl index 1391e6aef..b4913447c 100644 --- a/src/rulesets/LinearAlgebra/factorization.jl +++ b/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) diff --git a/test/rulesets/LinearAlgebra/factorization.jl b/test/rulesets/LinearAlgebra/factorization.jl index 1dd1aeb37..49160d314 100644 --- a/test/rulesets/LinearAlgebra/factorization.jl +++ b/test/rulesets/LinearAlgebra/factorization.jl @@ -44,6 +44,19 @@ function rand_eigen(T::Type, n::Int) return V * Diagonal(λ) / V end +# m×n matrix whose LU factorization has an exactly zero last pivot: since |L[i, j]| < 1, +# partial pivoting keeps the row order, and dyadic entries and power-of-two pivots make +# the factorization exact +function lu_singular_last_pivot(m, n) + q = min(m, n) + L = tril(rand(-3:3, m, q) / 4, -1) + L[diagind(L)] .= 1 + U = triu(rand(-3:3, q, n)) + U[diagind(U)] = rand((-4, -2, -1, 1, 2, 4), q) + U[q, q] = 0 + return L * U +end + @testset "Factorizations" begin @testset "lu decomposition" begin n = 10 @@ -55,14 +68,30 @@ end test_frule(lu!, randn(T, m, n), pivot ⊢ NoTangent()) end + @testset "zero last pivot: lu!(A, $pivot) for size(A)=($m, $n)" for + pivot in (LU_ROW_MAXIMUM, LU_NO_PIVOT), + m in (7, 10, 13) + + A = lu_singular_last_pivot(m, n) + @test iszero(last(diag(lu(A, pivot; check=false).U))) + if m <= n + test_frule(lu!, A, pivot ⊢ NoTangent(); fkwargs=(; check=false)) + else # not differentiable + ΔA = rand_tangent(A) + _, ∂F = frule((ZeroTangent(), ΔA), lu!, A, pivot; check=false) + @test all(isnan, ∂F.factors) + end + end @testset "check=false passed to primal function" begin Asingular = zeros(n, n) ΔAsingular = rand_tangent(Asingular) @test_throws SingularException frule( (ZeroTangent(), copy(ΔAsingular)), lu!, copy(Asingular), LU_ROW_MAXIMUM ) - frule((ZeroTangent(), ΔAsingular), lu!, Asingular, LU_ROW_MAXIMUM; check=false) - @test true # above line would have errored if this was not working right + _, ∂F = frule( + (ZeroTangent(), ΔAsingular), lu!, Asingular, LU_ROW_MAXIMUM; check=false + ) + @test all(isnan, ∂F.factors) # not differentiable end end @testset "lu rrule" begin @@ -73,14 +102,28 @@ end test_rrule(lu, randn(T, m, n), pivot ⊢ NoTangent()) end + @testset "zero last pivot: lu(A, $pivot) for size(A)=($m, $n)" for + pivot in (LU_ROW_MAXIMUM, LU_NO_PIVOT), + m in (7, 10, 13) + + A = lu_singular_last_pivot(m, n) + @test iszero(last(diag(lu(A, pivot; check=false).U))) + if m <= n + test_rrule(lu, A, pivot ⊢ NoTangent(); fkwargs=(; check=false)) + else # not differentiable + F, back = rrule(lu, A, pivot; check=false) + _, ∂A, _ = back(Tangent{typeof(F)}(; factors=rand_tangent(F.factors))) + @test all(isnan, ∂A) + end + end @testset "check=false passed to primal function" begin Asingular = zeros(n, n) F = lu(Asingular, LU_ROW_MAXIMUM; check=false) - ΔF = Tangent{typeof(F)}(; U=rand_tangent(F.U), L=rand_tangent(F.L)) + ΔF = Tangent{typeof(F)}(; factors=rand_tangent(F.factors)) @test_throws SingularException rrule(lu, Asingular, LU_ROW_MAXIMUM) _, back = rrule(lu, Asingular, LU_ROW_MAXIMUM; check=false) - back(ΔF) - @test true # above line would have errored if this was not working right + _, ∂A, _ = back(ΔF) + @test all(isnan, ∂A) # not differentiable end end @testset "LU" begin From 141485553b383075b6a95c362cf1e81bedc8b3d8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Thu, 15 Oct 2026 15:41:47 +0200 Subject: [PATCH 07/12] Fix format Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- test/runtests.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/test/runtests.jl b/test/runtests.jl index 22791e3cd..1df03e998 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -33,7 +33,9 @@ function include_test(path) if isempty(ARGS) || any(occursin(a, path) for a in ARGS) println("Testing $path:") # print so TravisCI doesn't timeout due to no output @time Base.include(@__MODULE__(), path) do ex - INTERPRET_TESTSETS && Meta.isexpr(ex, :macrocall) && ex.args[1] == Symbol("@testset") || return ex + INTERPRET_TESTSETS && + Meta.isexpr(ex, :macrocall) && + ex.args[1] == Symbol("@testset") || return ex return :(@interpret (() -> $ex)()) # interpret testsets using JuliaInterpreter end else From a533818a67fdfafa6335eeb162063ab964e306d5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Thu, 15 Oct 2026 15:42:49 +0200 Subject: [PATCH 08/12] Apply batched suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- test/rulesets/LinearAlgebra/factorization.jl | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/test/rulesets/LinearAlgebra/factorization.jl b/test/rulesets/LinearAlgebra/factorization.jl index 49160d314..a108dd919 100644 --- a/test/rulesets/LinearAlgebra/factorization.jl +++ b/test/rulesets/LinearAlgebra/factorization.jl @@ -68,8 +68,9 @@ end test_frule(lu!, randn(T, m, n), pivot ⊢ NoTangent()) end - @testset "zero last pivot: lu!(A, $pivot) for size(A)=($m, $n)" for - pivot in (LU_ROW_MAXIMUM, LU_NO_PIVOT), + @testset "zero last pivot: lu!(A, $pivot) for size(A)=($m, $n)" for pivot in ( + LU_ROW_MAXIMUM, LU_NO_PIVOT + ), m in (7, 10, 13) A = lu_singular_last_pivot(m, n) @@ -102,8 +103,9 @@ end test_rrule(lu, randn(T, m, n), pivot ⊢ NoTangent()) end - @testset "zero last pivot: lu(A, $pivot) for size(A)=($m, $n)" for - pivot in (LU_ROW_MAXIMUM, LU_NO_PIVOT), + @testset "zero last pivot: lu(A, $pivot) for size(A)=($m, $n)" for pivot in ( + LU_ROW_MAXIMUM, LU_NO_PIVOT + ), m in (7, 10, 13) A = lu_singular_last_pivot(m, n) From 7c4c8b4c56891f215c8a4f8d5eef34cf2cb77382 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Fri, 16 Oct 2026 00:11:25 +0200 Subject: [PATCH 09/12] Add `_sortslices_inds` helper and test `normalize` with subnormal norm - `sortslices` rules: factor out the thrice-repeated index tuple. - `normalize` tests: cover the subnormal-norm branch of `_normalize!`. Co-Authored-By: Claude Opus 5.5 --- src/rulesets/Base/sort.jl | 11 +++++++---- test/rulesets/LinearAlgebra/norm.jl | 8 ++++++++ 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/src/rulesets/Base/sort.jl b/src/rulesets/Base/sort.jl index de2bd4ece..88b02ccd8 100644 --- a/src/rulesets/Base/sort.jl +++ b/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...) + # the pullback does not capture the indices, whose type depends on the value of `dims` function sortslices_pullback(dy) - inds = ntuple(d -> d == dims ? p : (:), ndims(x)) - return (NoTangent(), ∇getindex(x, unthunk(dy), inds...)) + return (NoTangent(), ∇getindex(x, unthunk(dy), _sortslices_inds(x, dims, p)...)) end - return x[ntuple(d -> d == dims ? p : (:), ndims(x))...], 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` ##### diff --git a/test/rulesets/LinearAlgebra/norm.jl b/test/rulesets/LinearAlgebra/norm.jl index f43f5467e..2b6306966 100644 --- a/test/rulesets/LinearAlgebra/norm.jl +++ b/test/rulesets/LinearAlgebra/norm.jl @@ -199,6 +199,10 @@ end @gpu_broken test_rrule(normalize, rand(T, 3, 4)) test_rrule(normalize, adjoint(rand(T, 5))) + + # subnormal norm + xsub = fill(T(1e-310), 3) + @test rrule(normalize, xsub)[1] ≈ normalize(xsub) end @testset "x::Array{$T}, p=$p" for T in (Float64, ComplexF64), p in (1.0, 2.0, -Inf, Inf, 2.5) # skip p=0, since FD is unstable @@ -208,5 +212,9 @@ end @gpu_broken test_rrule(normalize, rand(T, 3, 4), p) test_rrule(normalize, adjoint(rand(T, 5)), p) + + # subnormal norm + xsub = fill(T(1e-310), 3) + @test rrule(normalize, xsub, p)[1] ≈ normalize(xsub, p) end end From 80106f115ed13bb7b3b39e50254d311b38ea7ebd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Fri, 16 Oct 2026 00:53:03 +0200 Subject: [PATCH 10/12] Test `Threads._spawn_set_thrpool` non-differentiable rule Co-Authored-By: Claude Opus 5.5 --- test/rulesets/Base/base.jl | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/test/rulesets/Base/base.jl b/test/rulesets/Base/base.jl index 5f16a8519..948fd585e 100644 --- a/test/rulesets/Base/base.jl +++ b/test/rulesets/Base/base.jl @@ -224,6 +224,13 @@ end @test rrule(Base.depwarn, "message", :f) !== nothing end + if isdefined(Threads, :_spawn_set_thrpool) + @testset "Threads" begin + @test frule((NoTangent(), NoTangent(), NoTangent()), Threads._spawn_set_thrpool, Task(() -> nothing), :default) !== nothing + @test rrule(Threads._spawn_set_thrpool, Task(() -> nothing), :default) !== nothing + 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)) From 44c4204a5ffd92c20b64426df165b45e5e1451e4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Fri, 16 Oct 2026 00:58:07 +0200 Subject: [PATCH 11/12] Test that `Threads._spawn_set_thrpool` rules return `NoTangent` Co-Authored-By: Claude Opus 5.5 --- test/rulesets/Base/base.jl | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/test/rulesets/Base/base.jl b/test/rulesets/Base/base.jl index 948fd585e..de803f4e2 100644 --- a/test/rulesets/Base/base.jl +++ b/test/rulesets/Base/base.jl @@ -226,8 +226,11 @@ end if isdefined(Threads, :_spawn_set_thrpool) @testset "Threads" begin - @test frule((NoTangent(), NoTangent(), NoTangent()), Threads._spawn_set_thrpool, Task(() -> nothing), :default) !== nothing - @test rrule(Threads._spawn_set_thrpool, Task(() -> nothing), :default) !== nothing + 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 From cb07681d32b057b1c87bc2a6d5d58282a4e2eb8f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Fri, 16 Oct 2026 00:59:37 +0200 Subject: [PATCH 12/12] Fix format Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- test/rulesets/Base/base.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/test/rulesets/Base/base.jl b/test/rulesets/Base/base.jl index de803f4e2..a71e7f603 100644 --- a/test/rulesets/Base/base.jl +++ b/test/rulesets/Base/base.jl @@ -227,7 +227,9 @@ end if isdefined(Threads, :_spawn_set_thrpool) @testset "Threads" begin f = Threads._spawn_set_thrpool - _, ∂y = frule((NoTangent(), NoTangent(), NoTangent()), f, Task(() -> nothing), :default) + _, ∂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())