[WIP] Route concatenation and zero! through TensorAlgebra

Project.toml

@@ -1,6 +1,6 @@
 name = "SparseArraysBase"
 uuid = "0d5efcca-f356-4864-8770-e1ed8d78f208"
-version = "0.10.7"
+version = "0.10.8"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]
@@ -17,12 +17,11 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MapBroadcast = "ebd9b9da-f48d-417c-9660-449667d60261"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
-
-[weakdeps]
 TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 
-[extensions]
-SparseArraysBaseTensorAlgebraExt = ["TensorAlgebra", "SparseArrays"]
+[sources.TensorAlgebra]
+rev = "mf/concatenate-directsum"
+url = "https://github.com/ITensor/TensorAlgebra.jl"
 
 [compat]
 Accessors = "0.1.41"
@@ -35,5 +34,5 @@ LinearAlgebra = "1.10"
 MapBroadcast = "0.1.5"
 Random = "1.10"
 SparseArrays = "1.10"
-TensorAlgebra = "0.15, 0.16, 0.17"
+TensorAlgebra = "0.17.5"
 julia = "1.10"

src/SparseArraysBase.jl

@@ -17,13 +17,10 @@ export SparseArrayDOK,
     storedpairs,
     storedvalues
 
-# `zero!` isn't defined in `Base`, but it is defined in `ArrayLayouts`
-# and is useful for sparse array logic, since it can be used to empty
-# the sparse array storage. SparseArraysBase owns its own `zero!` rather
-# than relying on an external definition.
-function zero! end
+# `zero!` is owned by TensorAlgebra; SparseArraysBase extends it so its sparse
+# types can empty their storage in place rather than filling with zeros.
+using TensorAlgebra: TensorAlgebra, zero!
 
-include("concatenate.jl")
 include("abstractsparsearraystyle.jl")
 include("sparsearraystyle.jl")
 include("indexing.jl")
@@ -33,5 +30,6 @@ include("abstractsparsearray.jl")
 include("sparsearraydok.jl")
 include("oneelementarray.jl")
 include("sparsearrays.jl")
+include("tensoralgebra.jl")
 
 end

src/abstractsparsearray.jl

@@ -90,7 +90,7 @@ Base.iszero(a::AnyAbstractSparseArray) = iszero_sparse(a)
 Base.isreal(a::AnyAbstractSparseArray) = isreal_sparse(a)
 Base.real(a::AnyAbstractSparseArray) = real_sparse(a)
 Base.fill!(a::AnyAbstractSparseArray, x) = fill!_sparse(a, x)
-zero!(a::AnyAbstractSparseArray) = zero!_sparse(a)
+TensorAlgebra.zero!(a::AnyAbstractSparseArray) = zero!_sparse(a)
 Base.zero(a::AnyAbstractSparseArray) = zero_sparse(a)
 function Base.permutedims!(dst, a::AnyAbstractSparseArray, perm)
     return permutedims!_sparse(dst, a, perm)
@@ -130,7 +130,7 @@ end
 using ArrayLayouts: ArrayLayouts
 ArrayLayouts.MemoryLayout(type::Type{<:AnyAbstractSparseArray}) = SparseLayout()
 
-using .Concatenate: concatenate
+using TensorAlgebra: concatenate
 # We overload `Base._cat` instead of `Base.cat` since it
 # is friendlier for invalidations/compile times, see:
 # https://github.com/ITensor/SparseArraysBase.jl/issues/25

src/sparsearraydok.jl

@@ -107,7 +107,7 @@ end
 storedpairs(a::SparseArrayDOK) = pairs(storage(a))
 
 # TODO: Also handle wrappers.
-function zero!(a::SparseArrayDOK)
+function TensorAlgebra.zero!(a::SparseArrayDOK)
     empty!(storage(a))
     return a
 end

src/sparsearraystyle.jl

@@ -2,13 +2,6 @@ function fill!_sparse(a::AbstractArray, value)
     return map!(Returns(value), a, a)
 end
 
-# Generic fallback for SparseArraysBase's owned `zero!` (see the declaration
-# in `SparseArraysBase.jl`).
-function zero!(a::AbstractArray)
-    fill!(a, zero(eltype(a)))
-    return a
-end
-
 # We use a single function definition to minimize method ambiguities.
 function zero!_sparse(a::AbstractArray)
     # More generally, this codepath could be taking if `zero(eltype(a))`

ext/SparseArraysBaseTensorAlgebraExt/SparseArraysBaseTensorAlgebraExt.jl → src/tensoralgebra.jl

@@ -1,9 +1,5 @@
-module SparseArraysBaseTensorAlgebraExt
-
 using SparseArrays: SparseMatrixCSC
-using SparseArraysBase: AnyAbstractSparseArray, AnyAbstractSparseMatrix, SparseArrayDOK
-using TensorAlgebra:
-    TensorAlgebra, FusionStyle, ReshapeFusion, bipermutedimsopadd!, matricize, unmatricize
+using TensorAlgebra: TensorAlgebra, FusionStyle, ReshapeFusion, matricize, unmatricize
 
 struct SparseArrayFusion <: FusionStyle end
 TensorAlgebra.FusionStyle(::Type{<:AnyAbstractSparseArray}) = SparseArrayFusion()
@@ -55,5 +51,3 @@ function _opadd!(dest::AbstractArray, op, src::AbstractArray, α, β)
     end
     return dest
 end
-
-end

test/Project.toml

@@ -34,5 +34,5 @@ SparseArrays = "1.10"
 SparseArraysBase = "0.10"
 StableRNGs = "1.0.2"
 Suppressor = "0.2.8"
-TensorAlgebra = "0.15, 0.16, 0.17"
+TensorAlgebra = "0.17.4"
 Test = "<0.0.1, 1"

test/test_tensoralgebraext.jl → test/test_tensoralgebra.jl

@@ -4,7 +4,7 @@ using SparseArraysBase:
 using TensorAlgebra: contract, matricizeperm
 using Test: @test, @testset
 
-@testset "TensorAlgebraExt (eltype = $elt)" for elt in (Float32, ComplexF64)
+@testset "TensorAlgebra (eltype = $elt)" for elt in (Float32, ComplexF64)
     a = sparsezeros(elt, (2, 2, 2))
     a[1, 1, 1] = 1
     a[2, 1, 2] = 2