diff --git a/Project.toml b/Project.toml index c5b6960..f8ba4c9 100644 --- a/Project.toml +++ b/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 and contributors"] [workspace] @@ -17,13 +17,8 @@ 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"] - [compat] Accessors = "0.1.41" Adapt = "4.3" @@ -35,5 +30,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" diff --git a/src/SparseArraysBase.jl b/src/SparseArraysBase.jl index 582a863..9c4fcc2 100644 --- a/src/SparseArraysBase.jl +++ b/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 diff --git a/src/abstractsparsearray.jl b/src/abstractsparsearray.jl index f1209ee..c46b0f6 100644 --- a/src/abstractsparsearray.jl +++ b/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 diff --git a/src/concatenate.jl b/src/concatenate.jl deleted file mode 100644 index 6fb92ca..0000000 --- a/src/concatenate.jl +++ /dev/null @@ -1,194 +0,0 @@ -# Alternative implementation for `Base.cat` through `Concatenate.cat(!)`. -# This is mostly a copy of the Base implementation, with the main difference being -# that the destination is chosen based on all inputs instead of just the first. -# There is an intermediate representation in terms of a `Concatenated` object, -# reminiscent of how Broadcast works. Destination selection can be customized through -# `Base.similar(::Concatenated{Style}, ::Type{T}, axes)`, and the operation itself -# through `Base.copy`/`Base.copyto!` on a `Concatenated`. -module Concatenate - -import Base.Broadcast as BC -using ..SparseArraysBase: zero! -using Base: promote_eltypeof - -unval(::Val{x}) where {x} = x - -function _Concatenated end - -# Lazy representation of the concatenation of various `Args` along `Dims`, in order to -# provide hooks to customize the implementation. -struct Concatenated{Style, Dims, Args <: Tuple} - style::Style - dims::Val{Dims} - args::Args - global @inline function _Concatenated( - style::Style, dims::Val{Dims}, args::Args - ) where {Style, Dims, Args <: Tuple} - return new{Style, Dims, Args}(style, dims, args) - end -end - -function Concatenated( - style::Union{BC.AbstractArrayStyle, Nothing}, dims::Val, args::Tuple - ) - return _Concatenated(style, dims, args) -end -function Concatenated(dims::Val, args::Tuple) - return Concatenated(cat_style(dims, args...), dims, args) -end -function Concatenated{Style}( - dims::Val, args::Tuple - ) where {Style <: Union{BC.AbstractArrayStyle, Nothing}} - return Concatenated(Style(), dims, args) -end - -dims(::Concatenated{<:Any, D}) where {D} = D -style(concat::Concatenated) = getfield(concat, :style) - -concatenated(dims, args...) = concatenated(Val(dims), args...) -concatenated(dims::Val, args...) = Concatenated(dims, args) - -function Base.convert( - ::Type{Concatenated{NewStyle}}, concat::Concatenated{<:Any, Dims, Args} - ) where {NewStyle, Dims, Args} - return Concatenated{NewStyle}( - concat.dims, concat.args - )::Concatenated{NewStyle, Dims, Args} -end - -# allocating the destination container -# ------------------------------------ -Base.similar(concat::Concatenated) = similar(concat, eltype(concat)) -Base.similar(concat::Concatenated, ::Type{T}) where {T} = similar(concat, T, axes(concat)) -function Base.similar(concat::Concatenated, ax) - return similar(concat, eltype(concat), ax) -end - -function Base.similar(concat::Concatenated, ::Type{T}, ax) where {T} - # Convert to a broadcasted to leverage its similar implementation. - bc = BC.Broadcasted(style(concat), identity, concat.args, ax) - return similar(bc, T) -end - -function cat_axis( - a1::AbstractUnitRange, a2::AbstractUnitRange, a_rest::AbstractUnitRange... - ) - return cat_axis(cat_axis(a1, a2), a_rest...) -end -function cat_axis(a1::AbstractUnitRange, a2::AbstractUnitRange) - first(a1) == first(a2) == 1 || throw(ArgumentError("Concatenated axes must start at 1")) - return Base.OneTo(length(a1) + length(a2)) -end - -function cat_ndims(dims, as::AbstractArray...) - return max(maximum(dims), maximum(ndims, as)) -end -function cat_ndims(dims::Val, as::AbstractArray...) - return cat_ndims(unval(dims), as...) -end - -function cat_axes(dims, a::AbstractArray, as::AbstractArray...) - return ntuple(cat_ndims(dims, a, as...)) do dim - return if dim in dims - cat_axis(map(Base.Fix2(axes, dim), (a, as...))...) - else - axes(a, dim) - end - end -end -function cat_axes(dims::Val, as::AbstractArray...) - return cat_axes(unval(dims), as...) -end - -function cat_style(dims, as::AbstractArray...) - N = cat_ndims(dims, as...) - return typeof(BC.combine_styles(as...))(Val(N)) -end - -Base.eltype(concat::Concatenated) = promote_eltypeof(concat.args...) -Base.axes(concat::Concatenated) = cat_axes(dims(concat), concat.args...) -Base.size(concat::Concatenated) = length.(axes(concat)) -Base.ndims(concat::Concatenated) = cat_ndims(dims(concat), concat.args...) - -# Main logic -# ---------- -# Concatenate the supplied `args` along dimensions `dims`. -concatenate(dims, args...) = Base.materialize(concatenated(dims, args...)) - -# Concatenate the supplied `args` along dimensions `dims`. -cat(args...; dims) = concatenate(dims, args...) -Base.materialize(concat::Concatenated) = copy(concat) - -# Concatenate the supplied `args` along dimensions `dims`, placing the result into `dest`. -function cat!(dest, args...; dims) - Base.materialize!(dest, concatenated(dims, args...)) - return dest -end -Base.materialize!(dest, concat::Concatenated) = copyto!(dest, concat) - -Base.copy(concat::Concatenated) = copyto!(similar(concat), concat) - -# The following is largely copied from the Base implementation of `Base.cat`, see: -# https://github.com/JuliaLang/julia/blob/885b1cd875f101f227b345f681cc36879124d80d/base/abstractarray.jl#L1778-L1887 -_copy_or_fill!(A, inds, x) = fill!(view(A, inds...), x) -_copy_or_fill!(A, inds, x::AbstractArray) = (A[inds...] = x) - -cat_size(A) = (1,) -cat_size(A::AbstractArray) = size(A) -cat_size(A, d) = 1 -cat_size(A::AbstractArray, d) = size(A, d) - -cat_indices(A, d) = Base.OneTo(1) -cat_indices(A::AbstractArray, d) = axes(A, d) - -function __cat!(A, shape, catdims, X...) - return __cat_offset!(A, shape, catdims, ntuple(zero, length(shape)), X...) -end -function __cat_offset!(A, shape, catdims, offsets, x, X...) - # splitting the "work" on x from X... may reduce latency (fewer costly specializations) - newoffsets = __cat_offset1!(A, shape, catdims, offsets, x) - return __cat_offset!(A, shape, catdims, newoffsets, X...) -end -__cat_offset!(A, shape, catdims, offsets) = A -function __cat_offset1!(A, shape, catdims, offsets, x) - inds = ntuple(length(offsets)) do i - return if (i <= length(catdims) && catdims[i]) - offsets[i] .+ cat_indices(x, i) - else - 1:shape[i] - end - end - _copy_or_fill!(A, inds, x) - newoffsets = ntuple(length(offsets)) do i - return if (i <= length(catdims) && catdims[i]) - offsets[i] + cat_size(x, i) - else - offsets[i] - end - end - return newoffsets -end - -dims2cat(dims::Val) = dims2cat(unval(dims)) -function dims2cat(dims) - if any(≤(0), dims) - throw(ArgumentError("All cat dimensions must be positive integers, but got $dims")) - end - return ntuple(in(dims), maximum(dims)) -end - -# default falls back to replacing style with Nothing -# this permits specializing on typeof(dest) without ambiguities -# Note: this needs to be defined for AbstractArray specifically to avoid ambiguities with Base. -@inline function Base.copyto!(dest::AbstractArray, concat::Concatenated) - return copyto!(dest, convert(Concatenated{Nothing}, concat)) -end - -function Base.copyto!(dest::AbstractArray, concat::Concatenated{Nothing}) - catdims = dims2cat(dims(concat)) - shape = size(concat) - count(!iszero, catdims)::Int > 1 && zero!(dest) - return __cat!(dest, shape, catdims, concat.args...) -end - -end diff --git a/src/sparsearraydok.jl b/src/sparsearraydok.jl index c08ca18..b1c53da 100644 --- a/src/sparsearraydok.jl +++ b/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 diff --git a/src/sparsearraystyle.jl b/src/sparsearraystyle.jl index 132fef0..a77e363 100644 --- a/src/sparsearraystyle.jl +++ b/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))` diff --git a/ext/SparseArraysBaseTensorAlgebraExt/SparseArraysBaseTensorAlgebraExt.jl b/src/tensoralgebra.jl similarity index 87% rename from ext/SparseArraysBaseTensorAlgebraExt/SparseArraysBaseTensorAlgebraExt.jl rename to src/tensoralgebra.jl index 7808033..6859aa5 100644 --- a/ext/SparseArraysBaseTensorAlgebraExt/SparseArraysBaseTensorAlgebraExt.jl +++ b/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 diff --git a/test/Project.toml b/test/Project.toml index b1efe2b..d8464e2 100644 --- a/test/Project.toml +++ b/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" diff --git a/test/test_tensoralgebraext.jl b/test/test_tensoralgebra.jl similarity index 92% rename from test/test_tensoralgebraext.jl rename to test/test_tensoralgebra.jl index 01543e5..f0e2c43 100644 --- a/test/test_tensoralgebraext.jl +++ b/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