diff --git a/Project.toml b/Project.toml index ba2eff7..2a54cf1 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "TensorAlgebra" uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" -version = "0.17.4" +version = "0.17.5" authors = ["ITensor developers and contributors"] [workspace] diff --git a/src/TensorAlgebra.jl b/src/TensorAlgebra.jl index a24d6c8..445d687 100644 --- a/src/TensorAlgebra.jl +++ b/src/TensorAlgebra.jl @@ -9,7 +9,7 @@ export contract, contract!, eig_full, eig_trunc, eig_vals, eigh_full, eigh_trunc if VERSION >= v"1.11.0-DEV.469" eval( Meta.parse( - "public biperm, bipartition, contractopadd!, data, datatype, flattenlinear, label_type, matricizeopperm, permutedims, permutedims!, scalar, similar_map, to_range, tr, tryflattenlinear, ungrade, zero!, scale!, permuteddims, PermutedDims" + "public biperm, bipartition, cat_similar, concatenate, concatenate!, contractopadd!, data, datatype, directsum, flattenlinear, label_type, matricizeopperm, permutedims, permutedims!, scalar, similar_map, to_range, tr, tryflattenlinear, ungrade, zero!, scale!, permuteddims, PermutedDims" ) ) end @@ -21,6 +21,8 @@ include("MatrixAlgebra.jl") include("bituple.jl") include("permutedimsadd.jl") include("matricize.jl") +include("concatenate.jl") +include("directsum.jl") include("diagonal.jl") include("to_range.jl") include("contract/contractalgorithm.jl") diff --git a/src/concatenate.jl b/src/concatenate.jl new file mode 100644 index 0000000..170d66c --- /dev/null +++ b/src/concatenate.jl @@ -0,0 +1,108 @@ +# `concatenate`/`concatenate!` over arrays, with the destination chosen from all inputs rather than +# just the first (unlike `Base.cat`). A backend customizes allocation by overloading `cat_similar` on +# the combined `cat_style` of the arguments, and placement by overloading `concatenate!` on the +# destination type. + +import Base.Broadcast as BC +using Base: promote_eltypeof + +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 + +cat_ndims(dims, as::AbstractArray...) = cat_ndims(Val(dims), as...) +function cat_ndims(dims::Val, as::AbstractArray...) + return max(maximum(unval(dims)), maximum(ndims, as)) +end + +cat_axes(dims, as::AbstractArray...) = cat_axes(Val(dims), as...) +function cat_axes(dims::Val, a::AbstractArray, as::AbstractArray...) + return ntuple(cat_ndims(dims, a, as...)) do dim + return if dim in unval(dims) + cat_axis(map(Base.Fix2(axes, dim), (a, as...))...) + else + axes(a, dim) + end + end +end + +function cat_style(dims, as::AbstractArray...) + N = cat_ndims(dims, as...) + return typeof(BC.combine_styles(as...))(Val(N)) +end + +# Default destination: reuse broadcast's `similar` for the style by wrapping in a `Broadcasted`. +function cat_similar(style, ::Type{T}, ax, args...) where {T} + return similar(BC.Broadcasted(style, identity, args, ax), T) +end + +concatenate(dims, args...) = concatenate(Val(dims), args...) +function concatenate(dims::Val, args...) + style = cat_style(dims, args...) + dest = cat_similar(style, promote_eltypeof(args...), cat_axes(dims, args...), args...) + return concatenate!(dest, dims, args...) +end + +# The offset placement below is adapted from Base's `cat`: +# 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) = dims2cat(Val(dims)) +function dims2cat(dims::Val) + d = unval(dims) + if any(≤(0), d) + throw(ArgumentError("All cat dimensions must be positive integers, but got $d")) + end + return ntuple(in(d), maximum(d)) +end + +# Generic placement; a backend overrides `concatenate!` on its destination type. +function concatenate!(dest, dims, args...) + catdims = dims2cat(dims) + shape = map(length, cat_axes(dims, args...)) + count(!iszero, catdims)::Int > 1 && zero!(dest) + return __cat!(dest, shape, catdims, args...) +end diff --git a/src/directsum.jl b/src/directsum.jl new file mode 100644 index 0000000..9a51eec --- /dev/null +++ b/src/directsum.jl @@ -0,0 +1,3 @@ +# `directsum` is a plain concatenation for now, kept as its own entry point so a fusing/rotating +# variant can later be selected by style, the way `matricize` takes a `FusionStyle`. +directsum(dims, as...) = concatenate(dims, as...) diff --git a/test/test_concatenate.jl b/test/test_concatenate.jl new file mode 100644 index 0000000..d0056f8 --- /dev/null +++ b/test/test_concatenate.jl @@ -0,0 +1,46 @@ +using TensorAlgebra: TensorAlgebra, cat_axes, concatenate, concatenate!, directsum +using Test: @test, @testset + +@testset "concatenate" begin + a = reshape(collect(1.0:6.0), 2, 3) + b = reshape(collect(7.0:12.0), 2, 3) + + # Single-dim concatenation matches `vcat`/`hcat`. + @test concatenate(1, a, b) == vcat(a, b) + @test concatenate(2, a, b) == hcat(a, b) + + # Multi-dim concatenation is the block-diagonal placement (off-diagonal blocks zeroed). + c = concatenate((1, 2), a, b) + ref = zeros(4, 6) + ref[1:2, 1:3] .= a + ref[3:4, 4:6] .= b + @test c == ref + + # `concatenate!` writes into a provided destination. + dest = zeros(4, 6) + @test concatenate!(dest, (1, 2), a, b) === dest + @test dest == ref + + # Element type is promoted across all inputs, not taken from the first. + @test eltype(concatenate(1, a, b .+ 0im)) == ComplexF64 + + # `cat_axes` computes the concatenated axes from the arguments. + @test cat_axes(Val((1, 2)), a, b) == (Base.OneTo(4), Base.OneTo(6)) +end + +@testset "directsum (forwards to concatenate)" begin + a = reshape(collect(1.0:6.0), 2, 3) + b = reshape(collect(7.0:12.0), 2, 3) + + # `directsum` is exactly `concatenate`: block-concatenation, no basis rotation. + @test directsum((1, 2), a, b) == concatenate((1, 2), a, b) + @test directsum(1, a, b) == vcat(a, b) + + # N-ary: each summand lands in its own diagonal hyper-block. + d = reshape(collect(1.0:8.0), 2, 4) + s = directsum((1, 2), a, b, d) + @test size(s) == (6, 10) + @test s[1:2, 1:3] == a + @test s[3:4, 4:6] == b + @test s[5:6, 7:10] == d +end diff --git a/test/test_exports.jl b/test/test_exports.jl index 7975506..183c3f5 100644 --- a/test/test_exports.jl +++ b/test/test_exports.jl @@ -38,9 +38,11 @@ using Test: @test, @testset append!( exports, [ - :biperm, :bipartition, :contractopadd!, :data, :datatype, :flattenlinear, - :label_type, :matricizeopperm, :permutedims, :permutedims!, :scalar, - :similar_map, :to_range, :tr, :tryflattenlinear, :ungrade, :zero!, :scale!, + :biperm, :bipartition, :cat_similar, + :concatenate, :concatenate!, :contractopadd!, :data, :datatype, :directsum, + :flattenlinear, :label_type, + :matricizeopperm, :permutedims, :permutedims!, :scalar, :similar_map, + :to_range, :tr, :tryflattenlinear, :ungrade, :zero!, :scale!, :permuteddims, :PermutedDims, ] )