Move SparseArrays, Statistics, Distributed and Random rules to extensions by devmotion · Pull Request #852 · JuliaDiff/ChainRules.jl

devmotion · 2026-10-22T23:08:50Z

SparseArrays, Statistics and Distributed are no longer part of the sysimage, but ChainRules loads all of them (SparseArrays also through SparseInverseSubset) for every package that depends on it. Their rules only apply once the respective package is loaded anyway, so this PR moves every rule set that can be disentangled to a package extension: SparseArraysExt, StatisticsExt, DistributedExt and RandomExt. LinearAlgebra, StructArrays, Adapt and GPUArraysCore stay hard dependencies since they are used throughout the other rules and the broadcasting machinery.

The PR consists of three commits, best reviewed separately:

Move SparseArrays, Statistics and Distributed rules to extensions. A pure move (detected as renames, so blame follows the files), with SparseInverseSubset still a dependency.
Replace SparseInverseSubset with a faster selected inverse. The rrules for det/logdet/logabsdet of SparseMatrixCSC only used SparseInverseSubset for the entries of the inverse on the pattern of the Cholesky factor. Its implementation scans all rows below the diagonal with sparse indexing for each entry. Computing Takahashi's recurrence directly on the pattern of the factor is about 100× faster (n = 10⁴, B'B + I with B = sprandn(n, n, 3/n): 23 s vs 2470 s, max. difference 5e-15) and removes the dependency. Additionally:
- factorize once in the forward pass and reuse the factorization in the pullback (previously logdet and det used an LU in the forward pass and a Cholesky factorization in the pullback)
- restrict the cotangent to the sparsity pattern of the input (previously the pattern of L + L')
- support matrices whose stored pattern is asymmetric due to explicit zeros (CHOLMOD's cholesky throws for these unless wrapped in Hermitian)
- the forward pass of logabsdet no longer throws for matrices that are not positive definite; pullbacks for non-positive-definite inputs throw an informative ArgumentError
- rules stay inferrable: the pullback always captures a factorization of the same type (of an empty matrix if the input is not Hermitian) plus a success flag
Move Random rules to an extension. Random is currently in the sysimage and loaded at startup, so RandomExt is loaded together with ChainRules (checked on Julia 1.10 and 1.13, without using Random) and the change is not observable today; it avoids loading Random should that change.

Tests: comparisons of the selected inverse with the dense inverse (real SPD, complex Hermitian, explicit zeros, exact cancellation in the fill-in), test_rrule on a tridiagonal matrix with dense inverse, non-positive-definite inputs, and a check that using ChainRules does not load SparseArrays, Statistics or Distributed.

Notes:

For positive definite inputs, det is now computed as exp(logdet(F)) (as det of a CHOLMOD factorization does), with relative error growing with |logdet| (at most ~1.6e-13 before overflow).
Very old versions of dependencies allowed by the current lower compat bounds might still load SparseArrays transitively; this would only affect load time.
Version bumped to 1.75.0.

🤖 Generated with Claude Code

None of these packages is part of the sysimage anymore, so loading them unconditionally slows down loading ChainRules for all dependents. Their rules only apply once the respective package is loaded anyway. This commit only moves the rules; SparseInverseSubset (which loads SparseArrays) is still a dependency and is removed in the next commit. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The rrules for `det`, `logdet` and `logabsdet` of `SparseMatrixCSC` only used SparseInverseSubset for the entries of the inverse on the sparsity pattern of the Cholesky factor. Its implementation scans all rows below the diagonal with sparse indexing for each entry; computing Takahashi's recurrence directly on the pattern of the factor is about 100x faster for n = 10^4 and removes a dependency that loads SparseArrays. Additionally: - factorize once in the forward pass and reuse the factorization in the pullback (previously `logdet` and `det` used an LU in the forward pass and a Cholesky factorization in the pullback) - restrict the cotangent to the sparsity pattern of the input - support matrices whose stored pattern is asymmetric due to explicit zeros - do not throw in the forward pass of `logabsdet` for matrices that are not positive definite; pullbacks of non-positive-definite inputs throw an informative `ArgumentError` - test that loading ChainRules does not load SparseArrays, Statistics or Distributed Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The Random rules are self-contained as well. Random is currently part of the sysimage and loaded at startup, so the extension is loaded together with ChainRules and the change is not observable; it avoids loading Random should that change in the future. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

github-actions · 2026-10-22T23:09:47Z

+
+import ChainRulesCore: rrule
+
+function rrule(::typeof(sparse), I::AbstractVector, J::AbstractVector, V::AbstractVector, m, n, combine::typeof(+))


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

function rrule(::typeof(sparse), I::AbstractVector, J::AbstractVector, V::AbstractVector, m, n, combine::typeof(+))

function rrule(

::typeof(sparse),

I::AbstractVector,

J::AbstractVector,

V::AbstractVector,

m,

n,

combine::typeof(+),

)

github-actions · 2026-10-22T23:09:48Z

+
+function rrule(::typeof(sparse), I::AbstractVector, J::AbstractVector, V::AbstractVector, m, n, combine::typeof(+))
+    project_V = ProjectTo(V)
+


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

github-actions · 2026-10-22T23:09:48Z

+    function sparse_pullback(Ω̄)
+        ΔΩ = unthunk(Ω̄)
+        ΔV = project_V(ΔΩ[I .+ m .* (J .- 1)])
+        return NoTangent(), NoTangent(), NoTangent(), ΔV, NoTangent(), NoTangent(), NoTangent()


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

return NoTangent(), NoTangent(), NoTangent(), ΔV, NoTangent(), NoTangent(), NoTangent()

return NoTangent(),

NoTangent(), NoTangent(), ΔV, NoTangent(), NoTangent(),

NoTangent()

github-actions · 2026-10-22T23:09:48Z

+    return sparse(I, J, V, m, n, combine), sparse_pullback
+end
+
+function rrule(::Type{T}, A::AbstractMatrix) where T <: AbstractSparseMatrix


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

function rrule(::Type{T}, A::AbstractMatrix) where T <: AbstractSparseMatrix

function rrule(::Type{T}, A::AbstractMatrix) where {T<:AbstractSparseMatrix}

github-actions · 2026-10-22T23:09:48Z

+    return T(A), sparse_pullback
+end
+
+function rrule(::Type{T}, v::AbstractVector) where T <: AbstractSparseVector


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

function rrule(::Type{T}, v::AbstractVector) where T <: AbstractSparseVector

function rrule(::Type{T}, v::AbstractVector) where {T<:AbstractSparseVector}

codecov · 2026-10-22T23:18:00Z

Codecov Report

❌ Patch coverage is 97.24138% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 96.64%. Comparing base (ec4c127) to head (db985f8).

Files with missing lines	Patch %	Lines
ext/SparseArraysExt.jl	97.24%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #852      +/-   ##
==========================================
+ Coverage   96.57%   96.64%   +0.06%     
==========================================
  Files          28       28              
  Lines        3618     3692      +74     
==========================================
+ Hits         3494     3568      +74     
  Misses        124      124

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

On recent SparseArrays, CHOLMOD supports single precision, so the rrules for `det`, `logdet` and `logabsdet` of a `SparseMatrixCSC{Float32}` returned `Float32` primals (and single-precision gradients) whereas the primals, computed with UMFPACK, are `Float64`. Also reuse the unit vector for columns outside of the pattern of the Cholesky factor, and add tests for complex Hermitian and `Int`/`Float32` matrices. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

github-actions · 2026-10-23T13:02:11Z

+        success || throw(ArgumentError(
+            "rrules for `det`, `logdet` and `logabsdet` of a `SparseMatrixCSC` only support " *
+            "Hermitian positive definite matrices"
+        ))


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

success || throw(ArgumentError(

"rrules for `det`, `logdet` and `logabsdet` of a `SparseMatrixCSC` only support " *

"Hermitian positive definite matrices"

))

success || throw(

ArgumentError(

"rrules for `det`, `logdet` and `logabsdet` of a `SparseMatrixCSC` only support " *

"Hermitian positive definite matrices",

),

)

github-actions · 2026-10-23T13:02:11Z

+                end
+            end
+        end
+        return SparseMatrixCSC(size(x)..., copy(SparseArrays.getcolptr(x)), copy(xrows), vals)


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

return SparseMatrixCSC(size(x)..., copy(SparseArrays.getcolptr(x)), copy(xrows), vals)

return SparseMatrixCSC(

size(x)..., copy(SparseArrays.getcolptr(x)), copy(xrows), vals

)

github-actions · 2026-10-23T13:02:11Z

+        end
+        return Ω, det_pullback
+    end
+


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

github-actions · 2026-10-23T13:02:11Z

+function rrule(::typeof(spdiagm), m::Integer, n::Integer, kv::Pair{<:Integer,<:AbstractVector}...)
+


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

function rrule(::typeof(spdiagm), m::Integer, n::Integer, kv::Pair{<:Integer,<:AbstractVector}...)

function rrule(

::typeof(spdiagm), m::Integer, n::Integer, kv::Pair{<:Integer,<:AbstractVector}...

)

github-actions · 2026-10-23T13:02:11Z

+        A = sparse([1, 3, 4, 2, 3, 4, 1, 2, 3, 1, 2, 4], [1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4],
+                   [1.0, 1, 1, 1, 1, -1, 1, 1, 3, 1, -1, 3])


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

A = sparse([1, 3, 4, 2, 3, 4, 1, 2, 3, 1, 2, 4], [1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4],

[1.0, 1, 1, 1, 1, -1, 1, 1, 3, 1, -1, 3])

A = sparse(

[1, 3, 4, 2, 3, 4, 1, 2, 3, 1, 2, 4],

[1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4],

[1.0, 1, 1, 1, 1, -1, 1, 1, 3, 1, -1, 3],

)

…pe directly Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

devmotion and others added 3 commits October 23, 2026 00:43

github-actions Bot reviewed Oct 22, 2026

View reviewed changes

github-actions Bot reviewed Oct 23, 2026

View reviewed changes

Construct the matrix factorized by CHOLMOD with the target element ty…

db985f8

…pe directly Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Navigation Menu

Repository navigation

Move SparseArrays, Statistics, Distributed and Random rules to extensions - #852

Move SparseArrays, Statistics, Distributed and Random rules to extensions#852
devmotion wants to merge 5 commits into
mainfrom
dmw/stdlib-extensions

devmotion commented Oct 22, 2026

Uh oh!

github-actions Bot Oct 22, 2026

Uh oh!

github-actions Bot Oct 22, 2026

Uh oh!

github-actions Bot Oct 22, 2026

Uh oh!

github-actions Bot Oct 22, 2026

Uh oh!

github-actions Bot Oct 22, 2026

Uh oh!

codecov Bot commented Oct 22, 2026 •

edited

Loading

Uh oh!

github-actions Bot Oct 23, 2026

Uh oh!

github-actions Bot Oct 23, 2026

Uh oh!

github-actions Bot Oct 23, 2026

Uh oh!

github-actions Bot Oct 23, 2026

Uh oh!

github-actions Bot Oct 23, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant


		import ChainRulesCore: rrule

		function rrule(::typeof(sparse), I::AbstractVector, J::AbstractVector, V::AbstractVector, m, n, combine::typeof(+))


		function rrule(::typeof(sparse), I::AbstractVector, J::AbstractVector, V::AbstractVector, m, n, combine::typeof(+))
		project_V = ProjectTo(V)

	function rrule(::Type{T}, A::AbstractMatrix) where T <: AbstractSparseMatrix
	function rrule(::Type{T}, A::AbstractMatrix) where {T<:AbstractSparseMatrix}

	function rrule(::Type{T}, v::AbstractVector) where T <: AbstractSparseVector
	function rrule(::Type{T}, v::AbstractVector) where {T<:AbstractSparseVector}

		function rrule(::typeof(spdiagm), m::Integer, n::Integer, kv::Pair{<:Integer,<:AbstractVector}...)

		A = sparse([1, 3, 4, 2, 3, 4, 1, 2, 3, 1, 2, 4], [1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4],
		[1.0, 1, 1, 1, 1, -1, 1, 1, 3, 1, -1, 3])

Conversation

devmotion commented Oct 22, 2026

Uh oh!

github-actions Bot Oct 22, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Oct 22, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Oct 22, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Oct 22, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Oct 22, 2026

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented Oct 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

github-actions Bot Oct 23, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Oct 23, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Oct 23, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Oct 23, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Oct 23, 2026

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

codecov Bot commented Oct 22, 2026 •

edited

Loading