From 6e8258086961161b0b8f098c4736eaac9dd1fb24 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Wed, 24 Jun 2026 10:13:22 -0400 Subject: [PATCH 1/9] Replace the blocked-tuple types with a single BiTuple and drop the BlockArrays extension Replaces TensorAlgebra's blocked-tuple type family with a single concrete `BiTuple`, and removes the BlockArrays extension. Matricization, contraction, and factorizations only ever split a tensor's indices into two groups, a codomain and a domain, so the general N-block machinery (a family of `AbstractBlockTuple` / `BlockedTuple` / blocked-permutation types plus a BlockArrays-style `Block` / `mortar` / broadcast apparatus) was more than that use needs. It is all replaced by one two-field type: ```julia struct BiTuple{N1, N2, T1 <: NTuple{N1, Any}, T2 <: NTuple{N2, Any}} t1::T1 t2::T2 end ``` The builder names are unchanged (`tuplemortar`, `permmortar`, `blockedperm`, `blockedpermvcat`, `blockedperm_indexin`, `trivialbiperm`, `blockpermute`), backed by a small accessor surface (`blocks`, `firstblock`, `lastblock`, `blocklengths`). A biperm is a `BiTuple` of `Int`s whose permutation property the perm builders validate at construction, the way `permutedims(a, perm)` validates its `perm`, rather than encoding it in a dedicated type. Removing the BlockArrays-style apparatus leaves `TensorAlgebraBlockArraysExt` with nothing to provide, so the extension and the `BlockArrays` weak-dependency are dropped. The `BlockReshapeFusion` matricization style it defined for block-structured arrays moves to GradedArrays, its only consumer. This is a breaking release: the removed blocked-tuple and blocked-permutation types, the `Block*` family, the N-block and ellipsis `blockedpermvcat` forms, and the BlockArrays extension are all gone. Downstream packages that only build and destructure biperms keep working through the unchanged builder names. --- Project.toml | 5 +- .../TensorAlgebraBlockArraysExt.jl | 6 - .../blockarrays.jl | 99 ------- .../blockedtuple.jl | 20 -- .../TensorAlgebraMooncakeExt.jl | 8 +- src/TensorAlgebra.jl | 3 +- src/bituple.jl | 82 ++++++ src/blockedpermutation.jl | 218 -------------- src/blockedtuple.jl | 277 ------------------ src/contract/blockedperms.jl | 2 +- src/contract/contract.jl | 28 +- src/factorizations.jl | 50 ++-- src/matricize.jl | 27 +- src/matrixfunctions.jl | 4 +- test/test_basics.jl | 6 +- test/test_bituple.jl | 59 ++++ test/test_blockarrays_contract.jl | 110 ------- test/test_blockedpermutation.jl | 163 ----------- test/test_blockedtuple.jl | 148 ---------- test/test_mooncakeext.jl | 14 +- 20 files changed, 212 insertions(+), 1117 deletions(-) delete mode 100644 ext/TensorAlgebraBlockArraysExt/TensorAlgebraBlockArraysExt.jl delete mode 100644 ext/TensorAlgebraBlockArraysExt/blockarrays.jl delete mode 100644 ext/TensorAlgebraBlockArraysExt/blockedtuple.jl create mode 100644 src/bituple.jl delete mode 100644 src/blockedpermutation.jl delete mode 100644 src/blockedtuple.jl create mode 100644 test/test_bituple.jl delete mode 100644 test/test_blockarrays_contract.jl delete mode 100644 test/test_blockedpermutation.jl delete mode 100644 test/test_blockedtuple.jl diff --git a/Project.toml b/Project.toml index c804782b..633566af 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "TensorAlgebra" uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" -version = "0.10.0" +version = "0.11.0" authors = ["ITensor developers and contributors"] [workspace] @@ -15,19 +15,16 @@ StridedViews = "4db3bf67-4bd7-4b4e-b153-31dc3fb37143" TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6" [weakdeps] -BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2" [extensions] -TensorAlgebraBlockArraysExt = "BlockArrays" TensorAlgebraGPUArraysCoreExt = "GPUArraysCore" TensorAlgebraMooncakeExt = "Mooncake" TensorAlgebraTensorOperationsExt = "TensorOperations" [compat] -BlockArrays = "1.7.2" EllipsisNotation = "1.8" GPUArraysCore = "0.2" LinearAlgebra = "1.10" diff --git a/ext/TensorAlgebraBlockArraysExt/TensorAlgebraBlockArraysExt.jl b/ext/TensorAlgebraBlockArraysExt/TensorAlgebraBlockArraysExt.jl deleted file mode 100644 index 7920990c..00000000 --- a/ext/TensorAlgebraBlockArraysExt/TensorAlgebraBlockArraysExt.jl +++ /dev/null @@ -1,6 +0,0 @@ -module TensorAlgebraBlockArraysExt - -include("blockarrays.jl") -include("blockedtuple.jl") - -end diff --git a/ext/TensorAlgebraBlockArraysExt/blockarrays.jl b/ext/TensorAlgebraBlockArraysExt/blockarrays.jl deleted file mode 100644 index 9a464066..00000000 --- a/ext/TensorAlgebraBlockArraysExt/blockarrays.jl +++ /dev/null @@ -1,99 +0,0 @@ -using BlockArrays: AbstractBlockArray, AbstractBlockedUnitRange, Block, BlockedArray, - blockedrange, blocklength, blocks, eachblockaxes1, mortar -using TensorAlgebra: TensorAlgebra, AbstractBlockTuple, BlockedTuple, FusionStyle, - ReshapeFusion, matricize, matricize_axes, tensor_product_axis, unmatricize - -struct BlockReshapeFusion <: FusionStyle end -TensorAlgebra.FusionStyle(::Type{<:AbstractBlockArray}) = BlockReshapeFusion() - -function TensorAlgebra.trivial_axis( - style::BlockReshapeFusion, side::Val{:codomain}, a::AbstractArray, - axes_codomain::Tuple{Vararg{AbstractUnitRange}}, - axes_domain::Tuple{Vararg{AbstractUnitRange}} - ) - return blockedrange([1]) -end -function mortar_axis(axs) - all(isone ∘ first, axs) || - throw(ArgumentError("Only one-based axes are supported")) - return blockedrange(length.(axs)) -end -function TensorAlgebra.tensor_product_axis( - style::BlockReshapeFusion, side::Val{:codomain}, - r1::AbstractUnitRange, r2::AbstractUnitRange - ) - (isone(first(r1)) && isone(first(r2))) || - throw(ArgumentError("Only one-based axes are supported")) - blockaxpairs = Iterators.product(eachblockaxes1(r1), eachblockaxes1(r2)) - blockaxs = vec(map(splat(tensor_product_axis), blockaxpairs)) - return mortar_axis(blockaxs) -end -function TensorAlgebra.matricize( - style::BlockReshapeFusion, a::AbstractArray, ndims_codomain::Val - ) - ax = matricize_axes(style, a, ndims_codomain) - reshaped_blocks_a = reshape(blocks(a), blocklength.(ax)) - bs = map(reshaped_blocks_a) do b - return matricize(b, ndims_codomain) - end - return mortar(bs, ax) -end -using BlockArrays: blocklengths -function TensorAlgebra.unmatricize( - ::BlockReshapeFusion, m::AbstractMatrix, - axes_codomain::Tuple{Vararg{AbstractUnitRange}}, - axes_domain::Tuple{Vararg{AbstractUnitRange}} - ) - ax = (axes_codomain..., axes_domain...) - reshaped_blocks_m = reshape(blocks(m), blocklength.(ax)) - bs = map(CartesianIndices(reshaped_blocks_m)) do I - block_axes_I = BlockedTuple( - map(ntuple(identity, length(ax))) do i - return Base.axes1(ax[i][Block(I[i])]) - end, - (length(axes_codomain), length(axes_domain)) - ) - return unmatricize(reshaped_blocks_m[I], block_axes_I) - end - return mortar(bs, ax) -end - -TensorAlgebra.FusionStyle(::Type{<:BlockedArray}) = ReshapeFusion() -unblock(a::BlockedArray) = a.blocks -unblock(a::AbstractBlockArray) = a[Base.OneTo.(size(a))...] -unblock(a::AbstractArray) = a -function TensorAlgebra.matricize(::ReshapeFusion, a::BlockedArray, ndims_codomain::Val) - return matricize(ReshapeFusion(), unblock(a), ndims_codomain) -end -function unmatricize_blocked( - style::ReshapeFusion, m::AbstractMatrix, - axes_codomain::Tuple{Vararg{AbstractUnitRange}}, - axes_domain::Tuple{Vararg{AbstractUnitRange}} - ) - a = unmatricize( - ReshapeFusion(), m, - Base.OneTo.(length.(axes_codomain)), Base.OneTo.(length.(axes_domain)) - ) - return BlockedArray(a, (axes_codomain..., axes_domain...)) -end -function TensorAlgebra.unmatricize( - style::ReshapeFusion, m::AbstractMatrix, - axes_codomain::Tuple{AbstractBlockedUnitRange, Vararg{AbstractBlockedUnitRange}}, - axes_domain::Tuple{AbstractBlockedUnitRange, Vararg{AbstractBlockedUnitRange}} - ) - return unmatricize_blocked(style, m, axes_codomain, axes_domain) -end -function TensorAlgebra.unmatricize( - style::ReshapeFusion, m::AbstractMatrix, - axes_codomain::Tuple{AbstractBlockedUnitRange, Vararg{AbstractBlockedUnitRange}}, - axes_domain::Tuple{Vararg{AbstractBlockedUnitRange}} - ) - return unmatricize_blocked(style, m, axes_codomain, axes_domain) -end -function TensorAlgebra.unmatricize( - style::ReshapeFusion, m::AbstractMatrix, - axes_codomain::Tuple{Vararg{AbstractBlockedUnitRange}}, - axes_domain::Tuple{AbstractBlockedUnitRange, Vararg{AbstractBlockedUnitRange}} - ) - return unmatricize_blocked(style, m, axes_codomain, axes_domain) -end diff --git a/ext/TensorAlgebraBlockArraysExt/blockedtuple.jl b/ext/TensorAlgebraBlockArraysExt/blockedtuple.jl deleted file mode 100644 index 84dd0a77..00000000 --- a/ext/TensorAlgebraBlockArraysExt/blockedtuple.jl +++ /dev/null @@ -1,20 +0,0 @@ -import BlockArrays as BA -import TensorAlgebra as TA - -BA.blockfirsts(bt::TA.AbstractBlockTuple) = TA.blockfirsts(bt) -BA.blocklasts(bt::TA.AbstractBlockTuple) = TA.blocklasts(bt) -BA.blocklength(bt::TA.AbstractBlockTuple) = TA.blocklength(bt) -BA.blocklengths(bt::TA.AbstractBlockTuple) = TA.blocklengths(bt) -BA.blocklengths(type::Type{<:TA.AbstractBlockTuple}) = TA.blocklengths(type) -BA.blocks(bt::TA.AbstractBlockTuple) = TA.blocks(bt) - -TA.Block(I::BA.Block) = TA.Block(I.n) -TA.BlockRange(I::BA.BlockRange) = TA.BlockRange(I.indices) -TA.BlockIndexRange(I::BA.BlockIndexRange) = TA.BlockIndexRange(TA.Block(I.block), I.indices) -Base.:(==)(I::BA.Block, J::TA.Block) = I.n == J.n -Base.:(==)(I::TA.Block, J::BA.Block) = I.n == J.n -Base.getindex(bt::TA.AbstractBlockTuple, I::BA.Block) = bt[TA.Block(I)] -Base.getindex(bt::TA.AbstractBlockTuple, I::BA.BlockIndexRange) = bt[TA.BlockIndexRange(I)] -Base.getindex(bt::TA.AbstractBlockTuple, I::BA.BlockRange{1}) = bt[TA.BlockRange(I)] - -BA.blocklasts(r::TA.BlockedOneTo) = TA.blocklasts(r) diff --git a/ext/TensorAlgebraMooncakeExt/TensorAlgebraMooncakeExt.jl b/ext/TensorAlgebraMooncakeExt/TensorAlgebraMooncakeExt.jl index 881160f2..9db0d909 100644 --- a/ext/TensorAlgebraMooncakeExt/TensorAlgebraMooncakeExt.jl +++ b/ext/TensorAlgebraMooncakeExt/TensorAlgebraMooncakeExt.jl @@ -1,11 +1,11 @@ module TensorAlgebraMooncakeExt using Mooncake: Mooncake, @zero_derivative, DefaultCtx -using TensorAlgebra: AbstractBlockPermutation, ContractAlgorithm, allocate_output, biperm, - blockedperms, check_input, contract, contract!, contract_labels, - default_contract_algorithm, select_contract_algorithm +using TensorAlgebra: BiTuple, ContractAlgorithm, allocate_output, biperm, blockedperms, + check_input, contract, contract!, contract_labels, default_contract_algorithm, + select_contract_algorithm -Mooncake.tangent_type(::Type{<:AbstractBlockPermutation}) = Mooncake.NoTangent +Mooncake.tangent_type(::Type{<:BiTuple}) = Mooncake.NoTangent Mooncake.tangent_type(::Type{<:ContractAlgorithm}) = Mooncake.NoTangent @zero_derivative DefaultCtx Tuple{ diff --git a/src/TensorAlgebra.jl b/src/TensorAlgebra.jl index 303067da..88d45400 100644 --- a/src/TensorAlgebra.jl +++ b/src/TensorAlgebra.jl @@ -11,8 +11,7 @@ end include("inplace.jl") include("MatrixAlgebra.jl") -include("blockedtuple.jl") -include("blockedpermutation.jl") +include("bituple.jl") include("BaseExtensions/BaseExtensions.jl") include("permutedimsadd.jl") include("matricize.jl") diff --git a/src/bituple.jl b/src/bituple.jl new file mode 100644 index 00000000..317227aa --- /dev/null +++ b/src/bituple.jl @@ -0,0 +1,82 @@ +# A two-block tuple: a flat tuple split into a codomain group `t1` and a domain group `t2`. +# This is the only blocked-tuple shape matricization, contraction, and factorizations ever use. +# When its entries are `Int`s forming a permutation it acts as a "biperm"; the permutation +# property is validated at construction by the perm builders (like `permutedims(a, perm)`), +# not encoded in a separate type. + +unval(::Val{N}) where {N} = N + +struct BiTuple{N1, N2, T1 <: NTuple{N1, Any}, T2 <: NTuple{N2, Any}} + t1::T1 + t2::T2 +end + +# Accessors (a small BlockArrays-inspired surface). +blocks(bt::BiTuple) = (bt.t1, bt.t2) +firstblock(bt::BiTuple) = bt.t1 +lastblock(bt::BiTuple) = bt.t2 +blocklengths(::BiTuple{N1, N2}) where {N1, N2} = (N1, N2) +blocklength(::BiTuple) = 2 + +Base.Tuple(bt::BiTuple) = (bt.t1..., bt.t2...) +Base.length(::BiTuple{N1, N2}) where {N1, N2} = N1 + N2 +Base.iterate(bt::BiTuple, args...) = iterate(Tuple(bt), args...) +Base.getindex(bt::BiTuple, i::Integer) = Tuple(bt)[i] +function Base.eltype(::Type{<:BiTuple{<:Any, <:Any, T1, T2}}) where {T1, T2} + return promote_type(eltype(T1), eltype(T2)) +end + +function Base.show(io::IO, bt::BiTuple) + return print(io, "tuplemortar(", blocks(bt), ")") +end + +function Base.invperm(bt::BiTuple{N1, N2}) where {N1, N2} + ip = invperm(Tuple(bt)) + return BiTuple(ntuple(i -> ip[i], Val(N1)), ntuple(i -> ip[N1 + i], Val(N2))) +end + +# +# Constructors +# + +# Axis bituple: split a tuple-of-tuples into the two groups verbatim. +tuplemortar(tt::Tuple{Tuple, Tuple}) = BiTuple(tt[1], tt[2]) + +# Permutation bituples: validate at runtime that the flat tuple is a permutation. +function permmortar(permblocks::Tuple{Tuple{Vararg{Int}}, Tuple{Vararg{Int}}}) + bt = BiTuple(permblocks[1], permblocks[2]) + @assert isperm(Tuple(bt)) + return bt +end + +# Split a flat permutation into a codomain block of length `blocklengths[1]` and the rest. +function blockedperm(perm::Tuple{Vararg{Int}}, blocklengths::Tuple{Int, Int}) + l1 = blocklengths[1] + return permmortar((perm[1:l1], perm[(l1 + 1):end])) +end + +# Two-block vcat builder (the contraction path's entry point). +function blockedpermvcat(block1::Tuple{Vararg{Int}}, block2::Tuple{Vararg{Int}}) + return permmortar((block1, block2)) +end + +# `indexin`-based builder: locate each group's labels within `collection`. +function blockedperm_indexin(collection, sub1, sub2) + return permmortar( + ( + BaseExtensions.indexin(sub1, collection), + BaseExtensions.indexin(sub2, collection), + ) + ) +end + +# Trivial (identity) biperm with the codomain/domain split given by `Val`s. Known to be a +# permutation by construction, so it skips the `isperm` check on this hot path. +function trivialbiperm(::Val{N1}, ::Val{N}) where {N1, N} + return BiTuple(ntuple(identity, Val(N1)), ntuple(i -> N1 + i, Val(N - N1))) +end + +# Bipartition a collection according to a biperm (out-of-place, blocked). +function blockpermute(v, bt::BiTuple) + return BiTuple(map(i -> v[i], bt.t1), map(i -> v[i], bt.t2)) +end diff --git a/src/blockedpermutation.jl b/src/blockedpermutation.jl deleted file mode 100644 index 982943af..00000000 --- a/src/blockedpermutation.jl +++ /dev/null @@ -1,218 +0,0 @@ -using EllipsisNotation: Ellipsis, var".." -using TupleTools: TupleTools - -trivialperm(len) = ntuple(identity, len) -function istrivialperm(t::Tuple) - return t == trivialperm(length(t)) -end - -unval(::Val{N}) where {N} = N - -_flatten_tuples(t::Tuple) = t -function _flatten_tuples(t1::Tuple, t2::Tuple, trest::Tuple...) - return _flatten_tuples((t1..., t2...), trest...) -end -_flatten_tuples() = () -flatten_tuples(ts::Tuple) = _flatten_tuples(ts...) - -collect_tuple(x) = (x,) -collect_tuple(x::Ellipsis) = x -collect_tuple(t::Tuple) = t - -# -# =============================== AbstractBlockPermutation =============================== -# -abstract type AbstractBlockPermutation{BlockLength} <: AbstractBlockTuple{BlockLength} end - -widened_constructorof(::Type{<:AbstractBlockPermutation}) = BlockedTuple - -# Otherwise it will convert to a BlockTuple since the default `Base.deepcopy` implementation -# calls `map`. -Base.deepcopy(bp::AbstractBlockPermutation) = bp - -# Block a permutation based on the specified lengths. -# blockperm((4, 3, 2, 1), (2, 2)) == blockedperm((4, 3), (2, 1)) -# TODO: Optimize with StaticNumbers.jl or generated functions, see: -# https://discourse.julialang.org/t/avoiding-type-instability-when-slicing-a-tuple/38567 -function blockedperm(perm::Tuple{Vararg{Int}}, blocklengths::Tuple{Vararg{Int}}) - return blockedperm(BlockedTuple(perm, blocklengths)) -end - -function blockedperm(perm::Tuple{Vararg{Int}}, BlockLengths::Val) - return blockedperm(BlockedTuple(perm, BlockLengths)) -end - -function Base.invperm(bp::AbstractBlockPermutation) - # use Val to preserve compile time info - return blockedperm(invperm(Tuple(bp)), Val(blocklengths(bp))) -end - -# interface - -# Bipartition a vector according to the -# bipartitioned permutation. -# Like `Base.permute!` block out-of-place and blocked. -function blockpermute(v, blockedperm::AbstractBlockPermutation) - return tuplemortar(map(blockperm -> map(i -> v[i], blockperm), blocks(blockedperm))) -end - -# -# Constructors -# - -function blockedperm(bt::AbstractBlockTuple) - return permmortar(blocks(bt)) -end - -# blockedpermvcat((4, 3), (2, 1)) -function blockedpermvcat( - permblocks::Tuple{Vararg{Int}}...; length::Union{Val, Nothing} = nothing - ) - return blockedpermvcat(length, permblocks...) -end - -function blockedpermvcat(::Nothing, permblocks::Tuple{Vararg{Int}}...) - return blockedpermvcat(Val(sum(length, permblocks; init = zero(Bool))), permblocks...) -end - -# blockedpermvcat((3, 2), 1) == blockedpermvcat((3, 2), (1,)) -function blockedpermvcat(permblocks::Union{Tuple{Vararg{Int}}, Int}...; kwargs...) - return blockedpermvcat(collect_tuple.(permblocks)...; kwargs...) -end - -function blockedpermvcat( - permblocks::Union{Tuple{Vararg{Int}}, Tuple{Ellipsis}, Int, Ellipsis}...; kwargs... - ) - return blockedpermvcat(collect_tuple.(permblocks)...; kwargs...) -end - -function blockedpermvcat(len::Val, permblocks::Tuple{Vararg{Int}}...) - unval(len) != sum(length.(permblocks); init = 0) && - throw(ArgumentError("Invalid total length")) - return permmortar(Tuple(permblocks)) -end - -function _blockedperm_length(::Nothing, specified_perm::Tuple{Vararg{Int}}) - return maximum(specified_perm) -end - -function _blockedperm_length(vallength::Val, ::Tuple{Vararg{Int}}) - return unval(vallength) -end - -# blockedpermvcat((4, 3), .., 1) == blockedpermvcat((4, 3), (2,), (1,)) -# blockedpermvcat((4, 3), .., 1; length=Val(5)) == blockedpermvcat((4, 3), (2,), (5,), (1,)) -# blockedpermvcat((4, 3), (..,), 1) == blockedpermvcat((4, 3), (2,), (1,)) -# blockedpermvcat((4, 3), (..,), 1; length=Val(5)) == blockedpermvcat((4, 3), (2, 5), (1,)) -function blockedpermvcat( - permblocks::Union{Tuple{Vararg{Int}}, Ellipsis, Tuple{Ellipsis}}...; - length::Union{Val, Nothing} = nothing - ) - return _blockedpermvcat(length, permblocks...) -end -function _blockedpermvcat( - length::Union{Val, Nothing}, - permblocks::Union{Tuple{Vararg{Int}}, Ellipsis, Tuple{Ellipsis}}... - ) - # Check there is only one `Ellipsis`. - @assert isone(count(x -> x isa Union{Ellipsis, Tuple{Ellipsis}}, permblocks)) - specified_permblocks = - filter(x -> !(x isa Union{Ellipsis, Tuple{Ellipsis}}), permblocks) - unspecified_dim = findfirst(x -> x isa Union{Ellipsis, Tuple{Ellipsis}}, permblocks) - specified_perm = flatten_tuples(specified_permblocks) - len = _blockedperm_length(length, specified_perm) - unspecified_dims_vec = setdiff(Base.OneTo(len), specified_perm) - ndims_unspecified = Val(len - sum(Base.length.(specified_permblocks))) # preserve type stability when possible - insert = unspecified_dims( - permblocks[unspecified_dim], unspecified_dims_vec, ndims_unspecified - ) - permblocks_specified = TupleTools.insertat(permblocks, unspecified_dim, insert) - return blockedpermvcat(permblocks_specified...) -end - -function unspecified_dims(::Tuple{Ellipsis}, unspecified_dims_vec, ndims_unspecified::Val) - return (ntuple(i -> unspecified_dims_vec[i], ndims_unspecified),) -end -function unspecified_dims(::Ellipsis, unspecified_dims_vec, ndims_unspecified::Val) - return ntuple(i -> (unspecified_dims_vec[i],), ndims_unspecified) -end - -# Version of `indexin` that outputs a `blockedperm`. -function blockedperm_indexin(collection, subs...) - return blockedpermvcat(map(sub -> BaseExtensions.indexin(sub, collection), subs)...) -end - -# -# ================================== BlockedPermutation ================================== -# - -# for dispatch reason, it is convenient to have BlockLength as the first parameter -struct BlockedPermutation{BlockLength, BlockLengths, Flat} <: - AbstractBlockPermutation{BlockLength} - flat::Flat - - function BlockedPermutation{BlockLength, BlockLengths}( - flat::Tuple - ) where {BlockLength, BlockLengths} - length(flat) != sum(BlockLengths; init = 0) && - throw(DimensionMismatch("Invalid total length")) - length(BlockLengths) != BlockLength && - throw(DimensionMismatch("Invalid total blocklength")) - any(BlockLengths .< 0) && throw(DimensionMismatch("Invalid block length")) - return new{BlockLength, BlockLengths, typeof(flat)}(flat) - end -end - -# Base interface -Base.Tuple(blockedperm::BlockedPermutation) = getfield(blockedperm, :flat) - -# BlockArrays interface -function blocklengths( - ::Type{<:BlockedPermutation{<:Any, BlockLengths}} - ) where {BlockLengths} - return BlockLengths -end - -function permmortar(permblocks::Tuple{Vararg{Tuple{Vararg{Int}}}}) - blockedperm = BlockedPermutation{length(permblocks), length.(permblocks)}( - flatten_tuples(permblocks) - ) - @assert isperm(blockedperm) - return blockedperm -end - -# -# ============================== BlockedTrivialPermutation =============================== -# -trivialperm(length::Union{Integer, Val}) = ntuple(identity, length) - -struct BlockedTrivialPermutation{BlockLength, BlockLengths} <: - AbstractBlockPermutation{BlockLength} end - -Base.Tuple(blockedperm::BlockedTrivialPermutation) = trivialperm(length(blockedperm)) - -# BlockArrays interface -function blocklengths( - ::Type{<:BlockedTrivialPermutation{<:Any, BlockLengths}} - ) where {BlockLengths} - return BlockLengths -end - -blockedperm(tp::BlockedTrivialPermutation) = tp - -function blockedtrivialperm(blocklengths::Tuple{Vararg{Val}}) - return BlockedTrivialPermutation{length(blocklengths), unval.(blocklengths)}() -end -function blockedtrivialperm(blocklengths::Tuple{Vararg{Int}}) - return blockedtrivialperm(Val.(blocklengths)) -end - -function trivialbiperm(length_codomain::Val, length::Val) - length_domain = Val(unval(length) - unval(length_codomain)) - return blockedtrivialperm((length_codomain, length_domain)) -end - -function trivialperm(blockedperm::AbstractBlockTuple) - return blockedtrivialperm(blocklengths(blockedperm)) -end -Base.invperm(blockedperm::BlockedTrivialPermutation) = blockedperm diff --git a/src/blockedtuple.jl b/src/blockedtuple.jl deleted file mode 100644 index ff69e23c..00000000 --- a/src/blockedtuple.jl +++ /dev/null @@ -1,277 +0,0 @@ -# This file defines an abstract type AbstractBlockTuple and a concrete type BlockedTuple. -# These types allow to store a Tuple of heterogeneous Tuples with a BlockArrays.jl like -# interface. - -# Local replacements for the two `TypeParameterAccessors` helpers this file used, so -# `TensorAlgebra` does not depend on that package. Same definitions it provided. -unspecify_type_parameters(type::Type) = Base.typename(type).wrapper -type_parameters(type::Type, position) = Base.unwrap_unionall(type).parameters[position] -type_parameters(object, position) = type_parameters(typeof(object), position) - -# BlockArrays-like interface -struct Block{N, T} - n::NTuple{N, T} - # Seems to be required for Aqua, see: - # https://github.com/JuliaTesting/Aqua.jl/issues/86 - function Block(n::Tuple{T, Vararg{T}}) where {T} - return new{length(n), T}(n) - end -end -Block(n::Integer) = Block((n,)) -Base.Int(b::Block{1}) = b.n[1] -Base.:(==)(I::Block, J::Block) = I.n == J.n -struct BlockIndex{N, TI <: Tuple{Vararg{Integer, N}}, Tα <: Tuple{Vararg{Any, N}}} - I::TI - α::Tα -end -struct BlockIndexRange{ - N, R <: Tuple{Vararg{AbstractUnitRange{<:Integer}, N}}, - I <: Tuple{Vararg{Any, N}}, BI, - } <: AbstractArray{BlockIndex{N, NTuple{N, BI}, I}, N} - block::Block{N, BI} - indices::R - function BlockIndexRange( - block::Block{N, BI}, inds::R - ) where {N, BI <: Integer, R <: Tuple{Vararg{AbstractUnitRange{<:Integer}, N}}} - I = Tuple{eltype.(inds)...} - return new{N, R, I, BI}(block, inds) - end -end -Block(I::BlockIndexRange) = I.block -struct BlockRange{N, R <: NTuple{N, AbstractUnitRange{<:Integer}}} <: - AbstractArray{Block{N, Int}, N} - indices::R -end -function Base.Broadcast.broadcasted( - ::Base.Broadcast.DefaultArrayStyle{1}, ::Type{<:Integer}, I::BlockRange{1} - ) - return first(I.indices) -end - -struct BlockedOneTo{T <: Integer, CS <: AbstractVector{T}} <: AbstractUnitRange{T} - lasts::CS -end -blockedrange(blocks) = BlockedOneTo(cumsum(blocks)) -function Base.length(a::BlockedOneTo) - return if isempty(blocklasts(a)) - zero(eltype(a)) - else - Integer(last(blocklasts(a)) - first(a) + oneunit(eltype(a))) - end -end -Base.first(b::BlockedOneTo) = oneunit(eltype(b)) -function Base.last(b::BlockedOneTo) - return isempty(blocklasts(b)) ? first(b) - oneunit(eltype(b)) : last(blocklasts(b)) -end -blocklasts(a::BlockedOneTo) = a.lasts - -# -# ================================== AbstractBlockTuple ================================== -# -# AbstractBlockTuple imposes BlockLength as first type parameter for easy dispatch -# it makes no assumption on storage type -abstract type AbstractBlockTuple{BlockLength} end - -constructorof(type::Type{<:AbstractBlockTuple}) = unspecify_type_parameters(type) -widened_constructorof(type::Type{<:AbstractBlockTuple}) = constructorof(type) - -# Like `BlockRange`. -function blockeachindex(bt::AbstractBlockTuple) - return ntuple(i -> Block(i), blocklength(bt)) -end - -# Base interface -Base.axes(bt::AbstractBlockTuple) = (blockedrange([blocklengths(bt)...]),) -Base.axes(::AbstractBlockTuple{0}) = (blockedrange(Int[]),) - -Base.deepcopy(bt::AbstractBlockTuple) = deepcopy.(bt) - -Base.firstindex(::AbstractBlockTuple) = 1 - -Base.getindex(bt::AbstractBlockTuple, i::Integer) = Tuple(bt)[i] -Base.getindex(bt::AbstractBlockTuple, r::AbstractUnitRange) = Tuple(bt)[r] -Base.getindex(bt::AbstractBlockTuple, b::Block{1}) = blocks(bt)[Int(b)] -function Base.getindex(bt::AbstractBlockTuple, br::BlockRange{1}) - r = Int.(br) - t = Tuple(bt)[blockfirsts(bt)[first(r)]:blocklasts(bt)[last(r)]] - return widened_constructorof(typeof(bt))(t, Val(blocklengths(bt)[r])) -end -function Base.getindex(bt::AbstractBlockTuple, bi::BlockIndexRange{1}) - return bt[Block(bi)][only(bi.indices)] -end -# needed for nested broadcast in Julia < 1.11 -Base.getindex(bt::AbstractBlockTuple, ci::CartesianIndex{1}) = bt[only(Tuple(ci))] - -Base.iterate(bt::AbstractBlockTuple) = iterate(Tuple(bt)) -Base.iterate(bt::AbstractBlockTuple, i::Int) = iterate(Tuple(bt), i) - -Base.lastindex(bt::AbstractBlockTuple) = length(bt) - -Base.length(bt::AbstractBlockTuple) = sum(blocklengths(bt); init = 0) - -function Base.map(f, bt::AbstractBlockTuple) - t = map(f, Tuple(bt)) - return widened_constructorof(typeof(bt))(t, Val(blocklengths(bt))) -end - -function Base.invperm(bt::AbstractBlockTuple) - return widened_constructorof(typeof(bt))(invperm(Tuple(bt)), Val(blocklengths(bt))) -end - -function Base.show(io::IO, bt::AbstractBlockTuple) - return print(io, nameof(typeof(bt)), blocks(bt)) -end -function Base.show(io::IO, ::MIME"text/plain", bt::AbstractBlockTuple) - println(io, typeof(bt)) - return print(io, blocks(bt)) -end - -# Broadcast interface -Base.broadcastable(bt::AbstractBlockTuple) = bt -struct AbstractBlockTupleBroadcastStyle{BlockLengths, BT} <: Broadcast.BroadcastStyle end -function Base.BroadcastStyle(T::Type{<:AbstractBlockTuple}) - return AbstractBlockTupleBroadcastStyle{blocklengths(T), unspecify_type_parameters(T)}() -end - -# default -combine_types(::Type{<:AbstractBlockTuple}, ::Type{<:AbstractBlockTuple}) = BlockedTuple - -# BroadcastStyle(::Style1, ::Style2) is not called when Style1 == Style2 -# tuplemortar(((1,), (2,))) .== tuplemortar(((1,), (2,))) = tuplemortar(((true,), (true,))) -# tuplemortar(((1,), (2,))) .== tuplemortar(((1, 2),)) = tuplemortar(((true,), (true,))) -# tuplemortar(((1,), (2,))) .== tuplemortar(((1,), (2,), (3,))) = error DimensionMismatch -function Base.BroadcastStyle( - s1::AbstractBlockTupleBroadcastStyle, s2::AbstractBlockTupleBroadcastStyle - ) - blocklengths1 = type_parameters(s1, 1) - blocklengths2 = type_parameters(s2, 1) - sum(blocklengths1; init = 0) != sum(blocklengths2; init = 0) && - throw(DimensionMismatch("blocked tuples could not be broadcast to a common size")) - new_blocklasts = static_mergesort(cumsum(blocklengths1), cumsum(blocklengths2)) - new_blocklengths = ( - first(new_blocklasts), - Base.tail(new_blocklasts) .- Base.front(new_blocklasts)..., - ) - BT = combine_types(type_parameters(s1, 2), type_parameters(s2, 2)) - return AbstractBlockTupleBroadcastStyle{new_blocklengths, BT}() -end - -static_mergesort(::Tuple{}, ::Tuple{}) = () -static_mergesort(a::Tuple, ::Tuple{}) = a -static_mergesort(::Tuple{}, b::Tuple) = b -function static_mergesort(a::Tuple, b::Tuple) - if first(a) == first(b) - return (first(a), static_mergesort(Base.tail(a), Base.tail(b))...) - end - if first(a) < first(b) - return (first(a), static_mergesort(Base.tail(a), b)...) - end - return (first(b), static_mergesort(a, Base.tail(b))...) -end - -# tuplemortar(((1,), (2,))) .== (1, 2) = (true, true) -function Base.BroadcastStyle( - s::AbstractBlockTupleBroadcastStyle, ::Base.Broadcast.Style{Tuple} - ) - return s -end - -# tuplemortar(((1,), (2,))) .== 1 = (true, false) -function Base.BroadcastStyle( - ::Base.Broadcast.DefaultArrayStyle{0}, s::AbstractBlockTupleBroadcastStyle - ) - return s -end - -# tuplemortar(((1,), (2,))) .== [1, 1] = BlockVector([true, false], [1, 1]) -function Base.BroadcastStyle( - a::Base.Broadcast.AbstractArrayStyle, ::AbstractBlockTupleBroadcastStyle - ) - return a -end - -function Base.copy( - bc::Broadcast.Broadcasted{AbstractBlockTupleBroadcastStyle{BlockLengths, BT}} - ) where {BlockLengths, BT} - t = bc.f.((Tuple.(bc.args))...) - return widened_constructorof(BT)(t, Val(BlockLengths)) -end - -Base.ndims(::Type{<:AbstractBlockTuple}) = 1 # needed in nested broadcast - -function Base.replace(bt::AbstractBlockTuple, replacements::Pair...; kwargs...) - t = replace(Tuple(bt), replacements...; kwargs...) - return widened_constructorof(typeof(bt))(t, Val(blocklengths(bt))) -end -# Copy of Base.Callable definition, since it isn't public. -const Callable = Union{Function, Type} -function Base.replace(f::Callable, bt::AbstractBlockTuple; kwargs...) - t = replace(f, Tuple(bt); kwargs...) - return widened_constructorof(typeof(bt))(t, Val(blocklengths(bt))) -end - -# BlockArrays interface -blockfirsts(::AbstractBlockTuple{0}) = () -function blockfirsts(bt::AbstractBlockTuple) - return (0, cumsum(Base.front(blocklengths(bt)))...) .+ 1 -end - -function blocklasts(bt::AbstractBlockTuple) - return cumsum(blocklengths(bt)) -end - -blocklength(::AbstractBlockTuple{BlockLength}) where {BlockLength} = BlockLength - -blocklengths(bt::AbstractBlockTuple) = blocklengths(typeof(bt)) - -function blocks(bt::AbstractBlockTuple) - bf = blockfirsts(bt) - bl = blocklasts(bt) - return ntuple(i -> Tuple(bt)[bf[i]:bl[i]], blocklength(bt)) -end - -# ===================================== BlockedTuple ===================================== -# -struct BlockedTuple{BlockLength, BlockLengths, Flat} <: AbstractBlockTuple{BlockLength} - flat::Flat - - function BlockedTuple{BlockLength, BlockLengths}( - flat::Tuple - ) where {BlockLength, BlockLengths} - length(BlockLengths) != BlockLength && - throw(DimensionMismatch("Invalid blocklength")) - length(flat) != sum(BlockLengths; init = 0) && - throw(DimensionMismatch("Invalid total length")) - any(BlockLengths .< 0) && throw(DimensionMismatch("Invalid block length")) - return new{BlockLength, BlockLengths, typeof(flat)}(flat) - end -end - -# TensorAlgebra Interface -function tuplemortar(tt::Tuple{Vararg{Tuple}}) - return BlockedTuple{length(tt), length.(tt)}(flatten_tuples(tt)) -end -function BlockedTuple(flat::Tuple, BlockLengths::Tuple{Vararg{Int}}) - return BlockedTuple{length(BlockLengths), BlockLengths}(flat) -end -function BlockedTuple(flat::Tuple, ::Val{BlockLengths}) where {BlockLengths} - # use Val to preserve compile time knowledge of BL - return BlockedTuple{length(BlockLengths), BlockLengths}(flat) -end -function BlockedTuple(bt::AbstractBlockTuple) - bl = blocklengths(bt) - return BlockedTuple{length(bl), bl}(Tuple(bt)) -end - -# Base interface -Base.Tuple(bt::BlockedTuple) = bt.flat - -# Forward the flat tuple's element type; otherwise `eltype` falls back to `Any`. -Base.eltype(::Type{<:BlockedTuple{<:Any, <:Any, Flat}}) where {Flat} = eltype(Flat) - -# BlockArrays interface -function blocklengths( - ::Type{<:BlockedTuple{<:Any, BlockLengths}} - ) where {BlockLengths} - return BlockLengths -end diff --git a/src/contract/blockedperms.jl b/src/contract/blockedperms.jl index f491e6e3..6615279f 100644 --- a/src/contract/blockedperms.jl +++ b/src/contract/blockedperms.jl @@ -9,7 +9,7 @@ function biperm(perm, ::Val{BlockLength1}) where {BlockLength1} return blockedperm(Tuple(perm), (BlockLength1, length(perm) - BlockLength1)) end -length_domain(t::AbstractBlockTuple{2}) = last(blocklengths(t)) +length_domain(t::BiTuple) = last(blocklengths(t)) # Assume all dimensions are in the codomain by default length_domain(t) = 0 diff --git a/src/contract/contract.jl b/src/contract/contract.jl index 6fe5815d..35b12904 100644 --- a/src/contract/contract.jl +++ b/src/contract/contract.jl @@ -171,16 +171,16 @@ end # BlockPermutation versions of contract[opadd][!] function contract( - a1::AbstractArray, biperm1::AbstractBlockPermutation{2}, - a2::AbstractArray, biperm2::AbstractBlockPermutation{2}; + a1::AbstractArray, biperm1::BiTuple, + a2::AbstractArray, biperm2::BiTuple; kwargs... ) return contract(a1, blocks(biperm1)..., a2, blocks(biperm2)...; kwargs...) end function contract( - biperm_dest::AbstractBlockPermutation{2}, - a1::AbstractArray, biperm1::AbstractBlockPermutation{2}, - a2::AbstractArray, biperm2::AbstractBlockPermutation{2}; + biperm_dest::BiTuple, + a1::AbstractArray, biperm1::BiTuple, + a2::AbstractArray, biperm2::BiTuple; kwargs... ) return contract( @@ -191,9 +191,9 @@ function contract( ) end function contract!( - a_dest::AbstractArray, biperm_dest::AbstractBlockPermutation{2}, - a1::AbstractArray, biperm1::AbstractBlockPermutation{2}, - a2::AbstractArray, biperm2::AbstractBlockPermutation{2}; + a_dest::AbstractArray, biperm_dest::BiTuple, + a1::AbstractArray, biperm1::BiTuple, + a2::AbstractArray, biperm2::BiTuple; kwargs... ) return contract!( @@ -204,9 +204,9 @@ function contract!( ) end function contractadd!( - a_dest::AbstractArray, biperm_dest::AbstractBlockPermutation{2}, - a1::AbstractArray, biperm1::AbstractBlockPermutation{2}, - a2::AbstractArray, biperm2::AbstractBlockPermutation{2}, + a_dest::AbstractArray, biperm_dest::BiTuple, + a1::AbstractArray, biperm1::BiTuple, + a2::AbstractArray, biperm2::BiTuple, α::Number, β::Number; kwargs... ) return contractadd!( @@ -217,9 +217,9 @@ function contractadd!( ) end function contractopadd!( - a_dest::AbstractArray, biperm_dest::AbstractBlockPermutation{2}, - op1, a1::AbstractArray, biperm1::AbstractBlockPermutation{2}, - op2, a2::AbstractArray, biperm2::AbstractBlockPermutation{2}, + a_dest::AbstractArray, biperm_dest::BiTuple, + op1, a1::AbstractArray, biperm1::BiTuple, + op2, a2::AbstractArray, biperm2::BiTuple, α::Number, β::Number; kwargs... ) return contractopadd!( diff --git a/src/factorizations.jl b/src/factorizations.jl index bee301db..3f3a1525 100644 --- a/src/factorizations.jl +++ b/src/factorizations.jl @@ -64,11 +64,11 @@ for f in ( function $f( style::FusionStyle, A::AbstractArray, - biperm::AbstractBlockPermutation{2}; kwargs... + biperm::BiTuple; kwargs... ) return $f(style, A, blocks(biperm)...; kwargs...) end - function $f(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) + function $f(A::AbstractArray, biperm::BiTuple; kwargs...) return $f(A, blocks(biperm)...; kwargs...) end end @@ -78,7 +78,7 @@ end qr_compact(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> Q, R qr_compact(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> Q, R qr_compact(A::AbstractArray, ndims_codomain::Val; kwargs...) -> Q, R - qr_compact(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> Q, R + qr_compact(A::AbstractArray, biperm::BiTuple; kwargs...) -> Q, R Compute the compact QR decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions, where `R` is square. The @@ -97,7 +97,7 @@ qr_compact qr_full(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> Q, R qr_full(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> Q, R qr_full(A::AbstractArray, ndims_codomain::Val; kwargs...) -> Q, R - qr_full(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> Q, R + qr_full(A::AbstractArray, biperm::BiTuple; kwargs...) -> Q, R Compute the full QR decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions, where `Q` is unitary. The @@ -116,7 +116,7 @@ qr_full lq_compact(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> L, Q lq_compact(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> L, Q lq_compact(A::AbstractArray, ndims_codomain::Val; kwargs...) -> L, Q - lq_compact(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> L, Q + lq_compact(A::AbstractArray, biperm::BiTuple; kwargs...) -> L, Q Compute the compact LQ decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions, where `L` is square. The @@ -135,7 +135,7 @@ lq_compact lq_full(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> L, Q lq_full(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> L, Q lq_full(A::AbstractArray, ndims_codomain::Val; kwargs...) -> L, Q - lq_full(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> L, Q + lq_full(A::AbstractArray, biperm::BiTuple; kwargs...) -> L, Q Compute the full LQ decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions, where `Q` is unitary. The @@ -154,7 +154,7 @@ lq_full left_polar(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> W, P left_polar(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> W, P left_polar(A::AbstractArray, ndims_codomain::Val; kwargs...) -> W, P - left_polar(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> W, P + left_polar(A::AbstractArray, biperm::BiTuple; kwargs...) -> W, P Compute the left polar decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. These can be specified either via @@ -172,7 +172,7 @@ left_polar right_polar(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> P, W right_polar(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> P, W right_polar(A::AbstractArray, ndims_codomain::Val; kwargs...) -> P, W - right_polar(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> P, W + right_polar(A::AbstractArray, biperm::BiTuple; kwargs...) -> P, W Compute the right polar decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. These can be specified either via @@ -190,7 +190,7 @@ right_polar left_orth(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> V, C left_orth(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> V, C left_orth(A::AbstractArray, ndims_codomain::Val; kwargs...) -> V, C - left_orth(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> V, C + left_orth(A::AbstractArray, biperm::BiTuple; kwargs...) -> V, C Compute the left orthogonal decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. These can be specified either via @@ -208,7 +208,7 @@ left_orth right_orth(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> C, V right_orth(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> C, V right_orth(A::AbstractArray, ndims_codomain::Val; kwargs...) -> C, V - right_orth(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> C, V + right_orth(A::AbstractArray, biperm::BiTuple; kwargs...) -> C, V Compute the right orthogonal decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. These can be specified either via @@ -279,7 +279,7 @@ end svd_compact(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> U, S, Vᴴ svd_compact(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> U, S, Vᴴ svd_compact(A::AbstractArray, ndims_codomain::Val; kwargs...) -> U, S, Vᴴ - svd_compact(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> U, S, Vᴴ + svd_compact(A::AbstractArray, biperm::BiTuple; kwargs...) -> U, S, Vᴴ Compute the compact (thin) SVD of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions, where `U` and `Vᴴ` are isometric. @@ -293,7 +293,7 @@ svd_compact svd_full(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> U, S, Vᴴ svd_full(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> U, S, Vᴴ svd_full(A::AbstractArray, ndims_codomain::Val; kwargs...) -> U, S, Vᴴ - svd_full(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> U, S, Vᴴ + svd_full(A::AbstractArray, biperm::BiTuple; kwargs...) -> U, S, Vᴴ Compute the full (thick) SVD of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions, where `U` and `Vᴴ` are unitary. @@ -307,7 +307,7 @@ svd_full svd_trunc(A::AbstractArray, labels_A, labels_codomain, labels_domain; trunc, kwargs...) -> U, S, Vᴴ svd_trunc(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; trunc, kwargs...) -> U, S, Vᴴ svd_trunc(A::AbstractArray, ndims_codomain::Val; trunc, kwargs...) -> U, S, Vᴴ - svd_trunc(A::AbstractArray, biperm::AbstractBlockPermutation{2}; trunc, kwargs...) -> U, S, Vᴴ + svd_trunc(A::AbstractArray, biperm::BiTuple; trunc, kwargs...) -> U, S, Vᴴ Compute the truncated SVD of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. The partition is specified either via @@ -326,7 +326,7 @@ svd_trunc svd_vals(A::AbstractArray, labels_A, labels_codomain, labels_domain) -> S svd_vals(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}) -> S svd_vals(A::AbstractArray, ndims_codomain::Val) -> S - svd_vals(A::AbstractArray, biperm::AbstractBlockPermutation{2}) -> S + svd_vals(A::AbstractArray, biperm::BiTuple) -> S Compute the singular values of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. The partition is specified either @@ -340,7 +340,7 @@ svd_vals eigh_full(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> D, V eigh_full(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> D, V eigh_full(A::AbstractArray, ndims_codomain::Val; kwargs...) -> D, V - eigh_full(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> D, V + eigh_full(A::AbstractArray, biperm::BiTuple; kwargs...) -> D, V Compute the eigenvalue decomposition of a generic N-dimensional array interpreted as a Hermitian linear map from the domain to the codomain dimensions. The partition is specified @@ -354,7 +354,7 @@ eigh_full eig_full(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> D, V eig_full(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> D, V eig_full(A::AbstractArray, ndims_codomain::Val; kwargs...) -> D, V - eig_full(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> D, V + eig_full(A::AbstractArray, biperm::BiTuple; kwargs...) -> D, V Compute the eigenvalue decomposition of a generic N-dimensional array interpreted as a general (non-Hermitian) linear map from the domain to the codomain dimensions. The output @@ -369,7 +369,7 @@ eig_full eigh_trunc(A::AbstractArray, labels_A, labels_codomain, labels_domain; trunc, kwargs...) -> D, V eigh_trunc(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; trunc, kwargs...) -> D, V eigh_trunc(A::AbstractArray, ndims_codomain::Val; trunc, kwargs...) -> D, V - eigh_trunc(A::AbstractArray, biperm::AbstractBlockPermutation{2}; trunc, kwargs...) -> D, V + eigh_trunc(A::AbstractArray, biperm::BiTuple; trunc, kwargs...) -> D, V Truncated Hermitian eigenvalue decomposition, like [`eigh_full`](@ref) but keeping only the eigenvalues selected by the `trunc` strategy. @@ -382,7 +382,7 @@ eigh_trunc eig_trunc(A::AbstractArray, labels_A, labels_codomain, labels_domain; trunc, kwargs...) -> D, V eig_trunc(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; trunc, kwargs...) -> D, V eig_trunc(A::AbstractArray, ndims_codomain::Val; trunc, kwargs...) -> D, V - eig_trunc(A::AbstractArray, biperm::AbstractBlockPermutation{2}; trunc, kwargs...) -> D, V + eig_trunc(A::AbstractArray, biperm::BiTuple; trunc, kwargs...) -> D, V Truncated general eigenvalue decomposition, like [`eig_full`](@ref) but keeping only the eigenvalues selected by the `trunc` strategy. @@ -395,7 +395,7 @@ eig_trunc eigh_vals(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> D eigh_vals(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> D eigh_vals(A::AbstractArray, ndims_codomain::Val; kwargs...) -> D - eigh_vals(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> D + eigh_vals(A::AbstractArray, biperm::BiTuple; kwargs...) -> D Compute the eigenvalues of a generic N-dimensional array interpreted as a Hermitian linear map from the domain to the codomain dimensions. The output is a vector of eigenvalues. @@ -408,7 +408,7 @@ eigh_vals eig_vals(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> D eig_vals(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> D eig_vals(A::AbstractArray, ndims_codomain::Val; kwargs...) -> D - eig_vals(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> D + eig_vals(A::AbstractArray, biperm::BiTuple; kwargs...) -> D Compute the eigenvalues of a generic N-dimensional array interpreted as a general (non-Hermitian) linear map from the domain to the codomain dimensions. The output is a @@ -422,7 +422,7 @@ eig_vals left_null(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> N left_null(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> N left_null(A::AbstractArray, ndims_codomain::Val; kwargs...) -> N - left_null(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> N + left_null(A::AbstractArray, biperm::BiTuple; kwargs...) -> N Compute the left nullspace of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. These can be specified either via @@ -462,7 +462,7 @@ end right_null(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> Nᴴ right_null(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> Nᴴ right_null(A::AbstractArray, ndims_codomain::Val::Val; kwargs...) -> Nᴴ - right_null(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> Nᴴ + right_null(A::AbstractArray, biperm::BiTuple; kwargs...) -> Nᴴ Compute the right nullspace of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. These can be specified either via @@ -502,7 +502,7 @@ end gram_eigh_full(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> X gram_eigh_full(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> X gram_eigh_full(A::AbstractArray, ndims_codomain::Val; kwargs...) -> X - gram_eigh_full(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> X + gram_eigh_full(A::AbstractArray, biperm::BiTuple; kwargs...) -> X Gram factorization of a generic N-dimensional array, interpreting it as a Hermitian positive semi-definite linear map from the domain to the codomain @@ -563,7 +563,7 @@ end gram_eigh_full_with_pinv(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> X, Y gram_eigh_full_with_pinv(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> X, Y gram_eigh_full_with_pinv(A::AbstractArray, ndims_codomain::Val; kwargs...) -> X, Y - gram_eigh_full_with_pinv(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> X, Y + gram_eigh_full_with_pinv(A::AbstractArray, biperm::BiTuple; kwargs...) -> X, Y Like [`gram_eigh_full`](@ref), but additionally returns `Y ≈ pinv(X)` such that `Y * X ≈ I` on the rank subspace (a left inverse). The codomain axes @@ -628,7 +628,7 @@ end TensorAlgebra.one(A::AbstractArray, labels_A, labels_codomain, labels_domain) -> Id TensorAlgebra.one(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}) -> Id TensorAlgebra.one(A::AbstractArray, ndims_codomain::Val) -> Id - TensorAlgebra.one(A::AbstractArray, biperm::AbstractBlockPermutation{2}) -> Id + TensorAlgebra.one(A::AbstractArray, biperm::BiTuple) -> Id Construct the identity operator tensor whose shape mirrors `A`, interpreted as a linear map from the domain to the codomain dimensions. The codomain and domain diff --git a/src/matricize.jl b/src/matricize.jl index 80f1925d..949e26a1 100644 --- a/src/matricize.jl +++ b/src/matricize.jl @@ -1,4 +1,5 @@ using .BaseExtensions: _permutedims, _permutedims! +using EllipsisNotation: Ellipsis using LinearAlgebra: Diagonal # ===================================== FusionStyle ====================================== @@ -138,7 +139,7 @@ end # Default similar with bipartitioned axes: flatten to a plain tuple of axes. # Downstream types (e.g., FusionTensor) can override to preserve bipartition. -function Base.similar(a::AbstractArray, T::Type, axes::BlockedTuple{2}) +function Base.similar(a::AbstractArray, T::Type, axes::BiTuple) return similar(a, T, Tuple(axes)) end @@ -171,11 +172,11 @@ end function bipermutedims!(a_dest::AbstractArray, a_src::AbstractArray, perm1, perm2) return bipermutedimsopadd!(a_dest, identity, a_src, perm1, perm2, true, false) end -function bipermutedims(a::AbstractArray, biperm::AbstractBlockPermutation{2}) +function bipermutedims(a::AbstractArray, biperm::BiTuple) return bipermutedims(a, blocks(biperm)...) end function bipermutedims!( - a_dest::AbstractArray, a_src::AbstractArray, biperm::AbstractBlockPermutation{2} + a_dest::AbstractArray, a_src::AbstractArray, biperm::BiTuple ) return bipermutedims!(a_dest, a_src, blocks(biperm)...) end @@ -245,11 +246,11 @@ function matricize( return matricize(style, a, to_permblocks(a, (perm_codomain, perm_domain))...) end -function matricize(a::AbstractArray, biperm_dest::AbstractBlockPermutation{2}) +function matricize(a::AbstractArray, biperm_dest::BiTuple) return matricize(FusionStyle(a), a, biperm_dest) end function matricize( - style::FusionStyle, a::AbstractArray, biperm_dest::AbstractBlockPermutation{2} + style::FusionStyle, a::AbstractArray, biperm_dest::BiTuple ) return matricize(style, a, blocks(biperm_dest)...) end @@ -304,11 +305,11 @@ function unmatricize( return unmatricize(FusionStyle(m), m, axes_codomain, axes_domain) end -function unmatricize(m::AbstractMatrix, blocked_axes::AbstractBlockTuple{2}) +function unmatricize(m::AbstractMatrix, blocked_axes::BiTuple) return unmatricize(FusionStyle(m), m, blocked_axes) end function unmatricize( - style::FusionStyle, m::AbstractMatrix, blocked_axes::AbstractBlockTuple{2} + style::FusionStyle, m::AbstractMatrix, blocked_axes::BiTuple ) return unmatricize(style, m, blocks(blocked_axes)...) end @@ -332,12 +333,12 @@ function unmatricize( return bipermutedims(a12, biperm_dest) end -function unmatricize(m::AbstractMatrix, axes_dest, invbiperm::AbstractBlockPermutation{2}) +function unmatricize(m::AbstractMatrix, axes_dest, invbiperm::BiTuple) return unmatricize(FusionStyle(m), m, axes_dest, invbiperm) end function unmatricize( style::FusionStyle, m::AbstractMatrix, axes_dest, - invbiperm::AbstractBlockPermutation{2} + invbiperm::BiTuple ) return unmatricize(style, m, axes_dest, blocks(invbiperm)...) end @@ -362,13 +363,13 @@ function unmatricize!( end function unmatricize!( - a_dest::AbstractArray, m::AbstractMatrix, invbiperm::AbstractBlockPermutation{2} + a_dest::AbstractArray, m::AbstractMatrix, invbiperm::BiTuple ) return unmatricize!(FusionStyle(m), a_dest, m, invbiperm) end function unmatricize!( style::FusionStyle, a_dest::AbstractArray, m::AbstractMatrix, - invbiperm::AbstractBlockPermutation{2} + invbiperm::BiTuple ) return unmatricize!(style, a_dest, m, blocks(invbiperm)...) end @@ -393,14 +394,14 @@ end function unmatricizeadd!( a_dest::AbstractArray, m::AbstractMatrix, - invbiperm::AbstractBlockPermutation{2}, + invbiperm::BiTuple, α::Number, β::Number ) return unmatricizeadd!(FusionStyle(a_dest), a_dest, m, invbiperm, α, β) end function unmatricizeadd!( style::FusionStyle, a_dest::AbstractArray, m::AbstractMatrix, - invbiperm::AbstractBlockPermutation{2}, + invbiperm::BiTuple, α::Number, β::Number ) return unmatricizeadd!( diff --git a/src/matrixfunctions.jl b/src/matrixfunctions.jl index 68b2674e..4d882fc8 100644 --- a/src/matrixfunctions.jl +++ b/src/matrixfunctions.jl @@ -79,11 +79,11 @@ for f in MATRIX_FUNCTIONS function $f( style::FusionStyle, a::AbstractArray, - biperm::AbstractBlockPermutation{2}; kwargs... + biperm::BiTuple; kwargs... ) return $f(style, a, blocks(biperm)...; kwargs...) end - function $f(a::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) + function $f(a::AbstractArray, biperm::BiTuple; kwargs...) return $f(a, blocks(biperm)...; kwargs...) end end diff --git a/test/test_basics.jl b/test/test_basics.jl index 096fe732..962d85f6 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -1,9 +1,9 @@ import TensorAlgebra using EllipsisNotation: var".." using StableRNGs: StableRNG -using TensorAlgebra: BlockedTuple, ContractAlgorithm, bipermutedims, bipermutedims!, - blockedpermvcat, contract, contract!, contractadd!, length_codomain, length_domain, - matricize, tuplemortar, unmatricize, unmatricize! +using TensorAlgebra: ContractAlgorithm, bipermutedims, bipermutedims!, blockedpermvcat, + contract, contract!, contractadd!, length_codomain, length_domain, matricize, + tuplemortar, unmatricize, unmatricize! using TensorOperations: TensorOperations using Test: @test, @test_broken, @test_throws, @testset diff --git a/test/test_bituple.jl b/test/test_bituple.jl new file mode 100644 index 00000000..4b842fa6 --- /dev/null +++ b/test/test_bituple.jl @@ -0,0 +1,59 @@ +using TensorAlgebra: BiTuple, blockedperm, blockedperm_indexin, blockedpermvcat, + blocklength, blocklengths, blockpermute, blocks, firstblock, lastblock, permmortar, + trivialbiperm, tuplemortar +using Test: @test, @test_throws, @testset +using TestExtras: @constinferred + +@testset "BiTuple (axis bituple)" begin + bt = @constinferred tuplemortar(((true, 'a'), (2.0,))) + @test bt isa BiTuple{2, 1} + @test (@constinferred Tuple(bt)) == (true, 'a', 2.0) + @test (@constinferred blocks(bt)) == ((true, 'a'), (2.0,)) + @test (@constinferred firstblock(bt)) == (true, 'a') + @test (@constinferred lastblock(bt)) == (2.0,) + @test (@constinferred blocklengths(bt)) == (2, 1) + @test blocklength(bt) == 2 + @test length(bt) == 3 + @test bt[1] == true + @test bt[3] == 2.0 + @test collect(bt) == [true, 'a', 2.0] + + bt_int = tuplemortar(((1,), (2, 3))) + @test eltype(bt_int) === Int + + # Empty blocks are allowed. + bt0 = @constinferred tuplemortar(((1,), ())) + @test blocks(bt0) == ((1,), ()) + @test blocklengths(bt0) == (1, 0) + @test Tuple(bt0) == (1,) +end + +@testset "BiTuple (biperm)" begin + p = @constinferred permmortar(((3, 4, 5), (2, 1))) + @test Tuple(p) === (3, 4, 5, 2, 1) + @test isperm(Tuple(p)) + @test blocks(p) == ((3, 4, 5), (2, 1)) + @test blocklengths(p) == (3, 2) + @test p == blockedpermvcat((3, 4, 5), (2, 1)) + @test p == blockedperm((3, 4, 5, 2, 1), (3, 2)) + @test Tuple(@constinferred invperm(p)) == invperm(Tuple(p)) + + # The perm builders validate that the flat tuple is a permutation. + @test_throws AssertionError permmortar(((3, 5), (2, 1))) + @test_throws AssertionError blockedpermvcat((0, 1), (2, 3)) + + # Trivial biperm: identity split into codomain/domain, built type-stably. + tb = @constinferred trivialbiperm(Val(2), Val(4)) + @test blocks(tb) == ((1, 2), (3, 4)) + @test Tuple(tb) == (1, 2, 3, 4) + @test blocks(@constinferred trivialbiperm(Val(0), Val(2))) == ((), (1, 2)) + + # Locate two label groups within a collection. + p = blockedperm_indexin(("a", "b", "c", "d"), ("c", "a"), ("b", "d")) + @test p == blockedpermvcat((3, 1), (2, 4)) + + # blockpermute splits a collection according to a biperm. + bp = blockpermute((10, 20, 30, 40), trivialbiperm(Val(1), Val(4))) + @test bp isa BiTuple + @test blocks(bp) == ((10,), (20, 30, 40)) +end diff --git a/test/test_blockarrays_contract.jl b/test/test_blockarrays_contract.jl deleted file mode 100644 index 7976ed3b..00000000 --- a/test/test_blockarrays_contract.jl +++ /dev/null @@ -1,110 +0,0 @@ -using BlockArrays: Block, BlockArray, BlockedArray, blockedrange, blocksize -using Random: randn! -using TensorAlgebra: contract, matricize, unmatricize -using Test: @test, @testset - -function randn_blockdiagonal(elt::Type, axes::Tuple) - a = zeros(elt, axes) - blockdiaglength = minimum(blocksize(a)) - for i in 1:blockdiaglength - b = Block(ntuple(Returns(i), ndims(a))) - a[b] = randn!(a[b]) - end - return a -end - -const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) -@testset "`contract` blocked arrays (eltype=$elt)" for elt in elts - d = blockedrange([2, 3]) - a1 = randn_blockdiagonal(elt, (d, d, d, d)) - a2 = randn_blockdiagonal(elt, (d, d, d, d)) - a3 = randn_blockdiagonal(elt, (d, d)) - a1_dense = convert(Array, a1) - a2_dense = convert(Array, a2) - a3_dense = convert(Array, a3) - - @testset "BlockedArray" begin - # matrix matrix - a_dest, dimnames_dest = contract(a1, (1, -1, 2, -2), a2, (2, -3, 1, -4)) - a_dest_dense, dimnames_dest_dense = contract( - a1_dense, (1, -1, 2, -2), a2_dense, (2, -3, 1, -4) - ) - @test dimnames_dest == dimnames_dest_dense - @test size(a_dest) == size(a_dest_dense) - @test a_dest isa BlockedArray{elt} - @test a_dest ≈ a_dest_dense - - # matrix vector - a_dest, dimnames_dest = contract(a1, (2, -1, -2, 1), a3, (1, 2)) - a_dest_dense, dimnames_dest_dense = - contract(a1_dense, (2, -1, -2, 1), a3_dense, (1, 2)) - @test dimnames_dest == dimnames_dest_dense - @test size(a_dest) == size(a_dest_dense) - @test a_dest isa BlockedArray{elt} - @test a_dest ≈ a_dest_dense - - # vector matrix - a_dest, dimnames_dest = contract(a3, (1, 2), a1, (2, -1, -2, 1)) - a_dest_dense, dimnames_dest_dense = - contract(a3_dense, (1, 2), a1_dense, (2, -1, -2, 1)) - @test dimnames_dest == dimnames_dest_dense - @test size(a_dest) == size(a_dest_dense) - @test a_dest isa BlockedArray{elt} - @test a_dest ≈ a_dest_dense - - # vector vector - a_dest, dimnames_dest = contract(a3, (1, 2), a3, (2, 1)) - a_dest_dense, dimnames_dest_dense = contract(a3_dense, (1, 2), a3_dense, (2, 1)) - @test dimnames_dest == dimnames_dest_dense - @test size(a_dest) == size(a_dest_dense) - @test a_dest isa BlockedArray{elt, 0} - @test a_dest ≈ a_dest_dense - - # outer product - a_dest, dimnames_dest = contract(a3, (1, 2), a3, (3, 4)) - a_dest_dense, dimnames_dest_dense = contract(a3_dense, (1, 2), a3_dense, (3, 4)) - @test dimnames_dest == dimnames_dest_dense - @test size(a_dest) == size(a_dest_dense) - @test a_dest isa BlockedArray{elt} - @test a_dest ≈ a_dest_dense - end - - @testset "BlockArray" begin - a1, a2, a3 = BlockArray.((a1, a2, a3)) - - # matrix matrix - a_dest, dimnames_dest = contract(a1, (1, -1, 2, -2), a2, (2, -3, 1, -4)) - m1 = matricize(a1, (2, 4), (1, 3)) - m2 = matricize(a2, (3, 1), (2, 4)) - m_dest = matricize(a_dest, Val(2)) - @test m_dest ≈ m1 * m2 - - # matrix vector - a_dest, dimnames_dest = contract(a1, (2, -1, -2, 1), a3, (1, 2)) - m1 = matricize(a1, (2, 3), (1, 4)) - m2 = matricize(a3, (2, 1), ()) - m_dest = matricize(a_dest, Val(2)) - @test m_dest ≈ m1 * m2 - - # vector matrix - a_dest, dimnames_dest = contract(a3, (1, 2), a1, (2, -1, -2, 1)) - m1 = matricize(a3, (), (1, 2)) - m2 = matricize(a1, (4, 1), (2, 3)) - m_dest = matricize(a_dest, Val(0)) - @test m_dest ≈ m1 * m2 - - # vector vector - a_dest, dimnames_dest = contract(a3, (1, 2), a3, (2, 1)) - m1 = matricize(a3, (), (1, 2)) - m2 = matricize(a3, (2, 1), ()) - m_dest = matricize(a_dest, Val(0)) - @test m_dest ≈ m1 * m2 - - # outer product - a_dest, dimnames_dest = contract(a3, (1, 2), a3, (3, 4)) - m1 = matricize(a3, (1, 2), ()) - m2 = matricize(a3, (), (1, 2)) - m_dest = matricize(a_dest, Val(2)) - @test m_dest ≈ m1 * m2 - end -end diff --git a/test/test_blockedpermutation.jl b/test/test_blockedpermutation.jl deleted file mode 100644 index 3be570a8..00000000 --- a/test/test_blockedpermutation.jl +++ /dev/null @@ -1,163 +0,0 @@ -using BlockArrays: blockfirsts, blocklasts, blocklength, blocklengths, blocks -using EllipsisNotation: var".." -using TensorAlgebra: BlockedPermutation, BlockedTrivialPermutation, BlockedTuple, - blockedperm, blockedperm_indexin, blockedpermvcat, blockedtrivialperm, blockpermute, - permmortar, trivialperm, tuplemortar -using Test: @test, @test_broken, @test_throws, @testset -using TestExtras: @constinferred - -@testset "BlockedPermutation" begin - p = @constinferred permmortar(((3, 4, 5), (2, 1))) - @test Tuple(p) === (3, 4, 5, 2, 1) - @test isperm(p) - @test length(p) == 5 - @test blocks(p) == ((3, 4, 5), (2, 1)) - @test blocklength(p) == 2 - @test blocklengths(p) == (3, 2) - @test blockfirsts(p) == (1, 4) - @test blocklasts(p) == (3, 5) - @test p == (@constinferred blockedpermvcat((3, 4, 5), (2, 1))) - @test p == blockedperm((3, 4, 5, 2, 1), (3, 2)) - @test p == (@constinferred blockedperm((3, 4, 5, 2, 1), Val((3, 2)))) - @test (@constinferred invperm(p)) == blockedpermvcat((5, 4, 1), (2, 3)) - @test p isa BlockedPermutation{2} - - flat = (3, 4, 5, 2, 1) - @test_throws DimensionMismatch BlockedPermutation{2, (1, 2, 2)}(flat) - @test_throws DimensionMismatch BlockedPermutation{3, (1, 2, 3)}(flat) - @test_throws DimensionMismatch BlockedPermutation{3, (-1, 3, 3)}(flat) - @test_throws AssertionError blockedpermvcat((3, 5), (2, 1)) - @test_throws AssertionError blockedpermvcat((0, 1), (2, 3)) - @test_throws AssertionError blockedpermvcat((0,)) - @test_throws AssertionError blockedpermvcat((2,)) - - # Empty block. - p = @constinferred blockedpermvcat((3, 2), (), (1,)) - @test Tuple(p) === (3, 2, 1) - @test isperm(p) - @test length(p) == 3 - @test blocks(p) == ((3, 2), (), (1,)) - @test blocklength(p) == 3 - @test blocklengths(p) == (2, 0, 1) - @test blockfirsts(p) == (1, 3, 3) - @test blocklasts(p) == (2, 2, 3) - @test invperm(p) == blockedpermvcat((3, 2), (), (1,)) - @test p isa BlockedPermutation{3} - - p = @constinferred blockedpermvcat((), ()) - @test Tuple(p) === () - @test blocklength(p) == 2 - @test blocklengths(p) == (0, 0) - @test isperm(p) - @test length(p) == 0 - @test blocks(p) == ((), ()) - @test p isa BlockedPermutation{2} - - p = @constinferred blockedpermvcat() - @test Tuple(p) === () - @test blocklength(p) == 0 - @test blocklengths(p) == () - @test isperm(p) - @test length(p) == 0 - @test blocks(p) == () - @test p isa BlockedPermutation{0} - - p = blockedpermvcat((3, 2), (), (1,)) - bt = tuplemortar(((3, 2), (), (1,))) - @test (@constinferred BlockedTuple(p)) == bt - @test (@constinferred map(identity, p)) == bt - @test (@constinferred p .+ p) == tuplemortar(((6, 4), (), (2,))) - @test (@constinferred p .+ bt) == tuplemortar(((6, 4), (), (2,))) - @test (@constinferred bt .+ p) == tuplemortar(((6, 4), (), (2,))) - @test (@constinferred blockedperm(p)) == p - @test (@constinferred blockedperm(bt)) == p - - @test_throws ArgumentError blockedpermvcat((1, 3), (2, 4); length = Val(6)) - - # Split collection into `BlockedPermutation`. - p = blockedperm_indexin(("a", "b", "c", "d"), ("c", "a"), ("b", "d")) - @test p == blockedpermvcat((3, 1), (2, 4)) - - # Singleton dimensions. - p = @constinferred blockedpermvcat((2, 3), 1) - @test p == blockedpermvcat((2, 3), (1,)) - - # First dimensions are unspecified. - p = blockedpermvcat(.., (4, 3)) - @test p == blockedpermvcat((1,), (2,), (4, 3)) - # Specify length - p = @constinferred blockedpermvcat(.., (4, 3); length = Val(6)) - @test p == blockedpermvcat((1,), (2,), (5,), (6,), (4, 3)) - - # Last dimensions are unspecified. - p = blockedpermvcat((4, 3), ..) - @test p == blockedpermvcat((4, 3), (1,), (2,)) - # Specify length - p = @constinferred blockedpermvcat((4, 3), ..; length = Val(6)) - @test p == blockedpermvcat((4, 3), (1,), (2,), (5,), (6,)) - - # Middle dimensions are unspecified. - p = blockedpermvcat((4, 3), .., 1) - @test p == blockedpermvcat((4, 3), (2,), (1,)) - # Specify length - # TODO: Type unstable in Julia 1.12, investigate. - # @constinferred blockedpermvcat((4, 3), .., 1; length = Val(6)) - p = blockedpermvcat((4, 3), .., 1; length = Val(6)) - @test p == blockedpermvcat((4, 3), (2,), (5,), (6,), (1,)) - - # No dimensions are unspecified. - p = blockedpermvcat((3, 2), .., 1) - @test p == blockedpermvcat((3, 2), (1,)) - - # same with (..,) instead of .. - p = blockedpermvcat((..,), (4, 3)) - @test p == blockedpermvcat((1, 2), (4, 3)) - p = @constinferred blockedpermvcat((..,), (4, 3); length = Val(6)) - @test p == blockedpermvcat((1, 2, 5, 6), (4, 3)) - - p = blockedpermvcat((4, 3), (..,)) - @test p == blockedpermvcat((4, 3), (1, 2)) - p = @constinferred blockedpermvcat((4, 3), (..,); length = Val(6)) - @test p == blockedpermvcat((4, 3), (1, 2, 5, 6)) - - p = blockedpermvcat((4, 3), (..,), 1) - @test p == blockedpermvcat((4, 3), (2,), (1,)) - # TODO: Type unstable in Julia 1.12, investigate. - # @constinferred blockedpermvcat((4, 3), (..,), 1; length = Val(6)) - p = blockedpermvcat((4, 3), (..,), 1; length = Val(6)) - @test p == blockedpermvcat((4, 3), (2, 5, 6), (1,)) - - p = blockedpermvcat((3, 2), (..,), 1) - @test p == blockedpermvcat((3, 2), (), (1,)) - - # blockpermute - t = (1, 2, 3, 4) - pblocks = tuplemortar(((4, 3), (), (1, 2))) - p = blockedperm(pblocks) - @test (@constinferred blockpermute(t, p)) isa - BlockedTuple{3, (2, 0, 2), NTuple{4, Int64}} - @test blockpermute(t, p) == pblocks - @test blockpermute(pblocks, p) == tuplemortar(((2, 1), (), (4, 3))) - @test blockpermute(p, p) == tuplemortar(((2, 1), (), (4, 3))) -end - -@testset "BlockedTrivialPermutation" begin - tp = blockedtrivialperm((2, 0, 1)) - - @test tp isa BlockedTrivialPermutation{3} - @test Tuple(tp) == (1, 2, 3) - @test blocklength(tp) == 3 - @test blocklengths(tp) == (2, 0, 1) - @test trivialperm(blockedpermvcat((3, 2), (), (1,))) == tp - - bt = tuplemortar(((1, 2), (), (3,))) - @test (@constinferred BlockedTuple(tp)) == bt - @test (@constinferred blocks(tp)) == blocks(bt) - @test (@constinferred map(identity, tp)) == bt - @test (@constinferred tp .+ tp) == tuplemortar(((2, 4), (), (6,))) - @test (@constinferred tp .+ Tuple(tp)) == tuplemortar(((2, 4), (), (6,))) - @test (@constinferred tp .+ BlockedTuple(tp)) == tuplemortar(((2, 4), (), (6,))) - @test (@constinferred blockedperm(tp)) == tp - @test (@constinferred trivialperm(tp)) == tp - @test (@constinferred trivialperm(bt)) == tp -end diff --git a/test/test_blockedtuple.jl b/test/test_blockedtuple.jl deleted file mode 100644 index c74dad59..00000000 --- a/test/test_blockedtuple.jl +++ /dev/null @@ -1,148 +0,0 @@ -using BlockArrays: - Block, BlockVector, blockedrange, blockisequal, blocklength, blocklengths, blocks -using TensorAlgebra: BlockedTuple, blockeachindex, tuplemortar -using Test: @test, @test_broken, @test_throws, @testset -using TestExtras: @constinferred - -@testset "BlockedTuple" begin - flat = (true, 'a', 2, "b", 3.0) - divs = (1, 2, 2) - - bt = @constinferred BlockedTuple{3, divs}(flat) - @test bt isa BlockedTuple{3} - @test (@constinferred blockeachindex(bt)) == (Block(1), Block(2), Block(3)) - - @test (@constinferred Tuple(bt)) == flat - @test (@constinferred tuplemortar(((true,), ('a', 2), ("b", 3.0)))) == bt - - # `eltype` forwards the underlying flat tuple's element type rather than `Any`. - @test eltype(bt) === eltype(flat) - bt_int = tuplemortar(((1,), (2, 3))) - @test eltype(bt_int) === Int - @test eltype(typeof(bt_int)) === Int - @test BlockedTuple(flat, divs) == bt - @test (@constinferred BlockedTuple(bt)) == bt - @test blocklength(bt) == 3 - @test blocklengths(bt) == (1, 2, 2) - @test (@constinferred blocks(bt)) == ((true,), ('a', 2), ("b", 3.0)) - - @test (@constinferred bt[1]) == true - @test (@constinferred bt[2]) == 'a' - @test (@constinferred map(identity, bt)) == bt - - # it is hard to make bt[Block(1)] type stable as compile-time knowledge of 1 is lost in Block - @test bt[Block(1)] == blocks(bt)[1] - @test bt[Block(2)] == blocks(bt)[2] - @test bt[Block(1):Block(2)] == tuplemortar(((true,), ('a', 2))) - @test bt[Block(2)[1:2]] == ('a', 2) - @test bt[2:4] == ('a', 2, "b") - - @test firstindex(bt) == 1 - @test lastindex(bt) == 5 - @test length(bt) == 5 - - @test iterate(bt) == (1, 2) - @test iterate(bt, 2) == ('a', 3) - @test blockisequal(only(axes(bt)), blockedrange([1, 2, 2])) - - @test_throws DimensionMismatch BlockedTuple{2, (1, 2, 2)}(flat) - @test_throws DimensionMismatch BlockedTuple{3, (1, 2, 3)}(flat) - @test_throws DimensionMismatch BlockedTuple{3, (-1, 3, 3)}(flat) - - bt = tuplemortar(((1,), (4, 2), (5, 3))) - @test bt isa BlockedTuple - @test Tuple(bt) == (1, 4, 2, 5, 3) - @test blocklengths(bt) == (1, 2, 2) - @test (@constinferred deepcopy(bt)) == bt - - @test (@constinferred map(n -> n + 1, bt)) == - BlockedTuple{3, blocklengths(bt)}(Tuple(bt) .+ 1) - @test (@constinferred bt .+ tuplemortar(((1,), (1, 1), (1, 1)))) == - BlockedTuple{3, blocklengths(bt)}(Tuple(bt) .+ 1) - @test (@constinferred bt .+ tuplemortar(((1,), (1, 1, 1), (1,)))) isa - BlockedTuple{4, (1, 2, 1, 1), NTuple{5, Int64}} - @test bt .+ tuplemortar(((1,), (1, 1, 1), (1,))) == - tuplemortar(((2,), (5, 3), (6,), (4,))) - - @test @constinferred(replace(bt, 4 => 40)) == - tuplemortar(((1,), (40, 2), (5, 3))) - @test @constinferred(replace(x -> x == 4 ? 40 : x, bt)) == - tuplemortar(((1,), (40, 2), (5, 3))) - - # `invperm` returns the inverse permutation as a `BlockedTuple` with the - # same block structure, so that callers that wrap permutations in a block - # layout (e.g. `LittleSet{BlockedTuple}` in NamedDimsArrays' broadcast - # alignment) can invert them transparently. - perm = tuplemortar(((3, 1), (2,))) - inv_perm = invperm(perm) - @test inv_perm isa BlockedTuple - @test Tuple(inv_perm) == invperm(Tuple(perm)) - @test blocklengths(inv_perm) == blocklengths(perm) - - bt = tuplemortar(((1:2, 1:2), (1:3,))) - @test length.(bt) == tuplemortar(((2, 2), (3,))) - # TODO: What is this testing? Should we fix it and bring it back? - @test_broken length.(length.(bt)) == tuplemortar(((1, 1), (1,))) - - bt = tuplemortar(((1,), (2,))) - @test (@constinferred bt .== bt) isa BlockedTuple{2, (1, 1), Tuple{Bool, Bool}} - @test (bt .== bt) == tuplemortar(((true,), (true,))) - @test (@constinferred bt .== tuplemortar(((1, 2),))) isa - BlockedTuple{2, (1, 1), Tuple{Bool, Bool}} - @test (bt .== tuplemortar(((1, 2),))) == tuplemortar(((true,), (true,))) - @test_throws DimensionMismatch bt .== tuplemortar(((1,), (2,), (3,))) - @test (@constinferred bt .== (1, 2)) isa BlockedTuple{2, (1, 1), Tuple{Bool, Bool}} - @test (bt .== (1, 2)) == tuplemortar(((true,), (true,))) - @test_throws DimensionMismatch bt .== (1, 2, 3) - @test (@constinferred bt .== 1) isa BlockedTuple{2, (1, 1), Tuple{Bool, Bool}} - @test (bt .== 1) == tuplemortar(((true,), (false,))) - @test (@constinferred bt .== (1,)) isa BlockedTuple{2, (1, 1), Tuple{Bool, Bool}} - - @test (bt .== (1,)) == tuplemortar(((true,), (false,))) - # BlockedTuple .== AbstractVector is not type stable. Requires fix in BlockArrays - # TODO: What is this testing? Should we fix it and bring it back? - @test_broken (bt .== [1, 1]) isa BlockVector{Bool} - # TODO: What is this testing? Should we fix it and bring it back? - @test_broken blocks(bt .== [1, 1]) == [[true], [false]] - @test_throws DimensionMismatch bt .== [1, 2, 3] - - @test (@constinferred (1, 2) .== bt) isa BlockedTuple{2, (1, 1), Tuple{Bool, Bool}} - @test ((1, 2) .== bt) == tuplemortar(((true,), (true,))) - @test_throws DimensionMismatch (1, 2, 3) .== bt - @test (@constinferred 1 .== bt) isa BlockedTuple{2, (1, 1), Tuple{Bool, Bool}} - @test (1 .== bt) == tuplemortar(((true,), (false,))) - @test (@constinferred (1,) .== bt) isa BlockedTuple{2, (1, 1), Tuple{Bool, Bool}} - @test ((1,) .== bt) == tuplemortar(((true,), (false,))) - # TODO: What is this testing? Should we fix it and bring it back? - @test_broken ([1, 1] .== bt) isa BlockVector{Bool} - # TODO: What is this testing? Should we fix it and bring it back? - @test_broken blocks([1, 1] .== bt) == [[true], [false]] - - # empty blocks - bt = tuplemortar(((1,), (), (5, 3))) - @test bt isa BlockedTuple{3} - @test Tuple(bt) == (1, 5, 3) - @test blocklengths(bt) == (1, 0, 2) - @test (@constinferred blocks(bt)) == ((1,), (), (5, 3)) - @test blockisequal(only(axes(bt)), blockedrange([1, 0, 2])) - - bt = tuplemortar(((), ())) - @test bt isa BlockedTuple{2} - @test Tuple(bt) == () - @test blocklengths(bt) == (0, 0) - @test (@constinferred blocks(bt)) == ((), ()) - @test blockisequal(only(axes(bt)), blockedrange([0, 0])) - @test bt == bt .+ bt - - bt0 = tuplemortar(()) - bt1 = tuplemortar(((),)) - @test bt0 isa BlockedTuple{0} - @test Tuple(bt0) == () - @test blocklengths(bt0) == () - @test (@constinferred blocks(bt0)) == () - @test blockisequal(only(axes(bt0)), blockedrange(zeros(Int, 0))) - @test bt0 == bt0 - @test bt != bt1 - @test (@constinferred bt0 .+ bt0) == bt0 - @test (@constinferred bt0 .+ bt1) == bt1 -end diff --git a/test/test_mooncakeext.jl b/test/test_mooncakeext.jl index 565d75ba..a7c9dcae 100644 --- a/test/test_mooncakeext.jl +++ b/test/test_mooncakeext.jl @@ -1,10 +1,9 @@ -using BlockArrays: blocks using Mooncake: Mooncake using Random: Random -using TensorAlgebra: AbstractBlockPermutation, BlockedPermutation, ContractAlgorithm, - DefaultContractAlgorithm, Matricize, allocate_output, biperm, blockedperms, check_input, - contract, contract!, contract_labels, contractadd!, default_contract_algorithm, - permmortar, select_contract_algorithm +using TensorAlgebra: BiTuple, ContractAlgorithm, DefaultContractAlgorithm, Matricize, + allocate_output, biperm, blockedperms, blocks, check_input, contract, contract!, + contract_labels, contractadd!, default_contract_algorithm, permmortar, + select_contract_algorithm using Test: @test, @testset @testset "MooncakeExt" begin @@ -15,8 +14,7 @@ using Test: @test, @testset atol = eps(real(elt))^(3 / 4) rtol = eps(real(elt))^(3 / 4) @testset "zero derivatives" begin - @test Mooncake.tangent_type(AbstractBlockPermutation) ≡ Mooncake.NoTangent - @test Mooncake.tangent_type(BlockedPermutation) ≡ Mooncake.NoTangent + @test Mooncake.tangent_type(BiTuple) ≡ Mooncake.NoTangent @test Mooncake.tangent_type(ContractAlgorithm) ≡ Mooncake.NoTangent @test Mooncake.tangent_type(DefaultContractAlgorithm) ≡ Mooncake.NoTangent @test Mooncake.tangent_type(Matricize) ≡ Mooncake.NoTangent @@ -65,7 +63,7 @@ using Test: @test, @testset @testset "contract" begin α = true β = false - @testset "contractadd! (BlockedPermutation)" begin + @testset "contractadd! (BiTuple)" begin dest = randn(elt, (2, 2)) a1 = randn(elt, (2, 2)) a2 = randn(elt, (2, 2)) From 5f413b3517a818a6edc7a5cfc292a0f931b6e04b Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Wed, 24 Jun 2026 12:31:39 -0400 Subject: [PATCH 2/9] Replace the blocked-tuple helpers with bipartition, biindexin, and tuple_indexin Collapses the remaining blocked-tuple interface to direct BiTuple field access plus a small set of plain-tuple helpers. bipartition splits a tuple into its two groups (optionally permuting), biindexin locates two label groups in a collection, and tuple_indexin is the tuple-preserving indexin. Drops trivialbiperm, tupleoneto, blockedperm_indexin, and the BaseExtensions submodule, whose remaining contents were either dead code or folded into tuple_indexin. --- src/BaseExtensions/BaseExtensions.jl | 4 -- src/BaseExtensions/indexin.jl | 5 -- src/BaseExtensions/permutedims.jl | 22 -------- src/TensorAlgebra.jl | 1 - src/bituple.jl | 81 ++++++++-------------------- src/contract/allocate_output.jl | 13 +++-- src/contract/blockedperms.jl | 35 +++++++----- src/contract/contract.jl | 41 +++++++------- src/contract/contract_matricize.jl | 4 +- src/factorizations.jl | 63 ++++++++++------------ src/matricize.jl | 35 +++++------- src/matrixfunctions.jl | 19 ++++--- test/test_basics.jl | 34 ++++++------ test/test_bituple.jl | 75 ++++++++++++-------------- test/test_mooncakeext.jl | 33 ++++++------ 15 files changed, 188 insertions(+), 277 deletions(-) delete mode 100644 src/BaseExtensions/BaseExtensions.jl delete mode 100644 src/BaseExtensions/indexin.jl delete mode 100644 src/BaseExtensions/permutedims.jl diff --git a/src/BaseExtensions/BaseExtensions.jl b/src/BaseExtensions/BaseExtensions.jl deleted file mode 100644 index c994fd81..00000000 --- a/src/BaseExtensions/BaseExtensions.jl +++ /dev/null @@ -1,4 +0,0 @@ -module BaseExtensions -include("indexin.jl") -include("permutedims.jl") -end diff --git a/src/BaseExtensions/indexin.jl b/src/BaseExtensions/indexin.jl deleted file mode 100644 index b25e2259..00000000 --- a/src/BaseExtensions/indexin.jl +++ /dev/null @@ -1,5 +0,0 @@ -# `Base.indexin` doesn't handle tuples -indexin(x, y::AbstractArray) = Base.indexin(x, y) -indexin(x, y) = Base.indexin(x, collect(y)) -indexin(x::Tuple, y::AbstractArray) = Tuple{Vararg{Any, length(x)}}(Base.indexin(x, y)) -indexin(x::Tuple, y) = Tuple{Vararg{Any, length(x)}}(Base.indexin(x, collect(y))) diff --git a/src/BaseExtensions/permutedims.jl b/src/BaseExtensions/permutedims.jl deleted file mode 100644 index c36e9409..00000000 --- a/src/BaseExtensions/permutedims.jl +++ /dev/null @@ -1,22 +0,0 @@ -# Workaround for https://github.com/JuliaLang/julia/issues/52615. -# Fixed by https://github.com/JuliaLang/julia/pull/52623. -# TODO remove once support for Julia 1.10 is dropped -function _permutedims!( - a_dest::AbstractArray{<:Any, N}, a_src::AbstractArray{<:Any, N}, - perm::Tuple{Vararg{Int, N}} - ) where {N} - permutedims!(a_dest, a_src, perm) - return a_dest -end -function _permutedims!( - a_dest::AbstractArray{<:Any, 0}, a_src::AbstractArray{<:Any, 0}, perm::Tuple{} - ) - a_dest[] = a_src[] - return a_dest -end -function _permutedims(a::AbstractArray{<:Any, N}, perm::Tuple{Vararg{Int, N}}) where {N} - return permutedims(a, perm) -end -function _permutedims(a::AbstractArray{<:Any, 0}, perm::Tuple{}) - return copy(a) -end diff --git a/src/TensorAlgebra.jl b/src/TensorAlgebra.jl index 88d45400..76fe137c 100644 --- a/src/TensorAlgebra.jl +++ b/src/TensorAlgebra.jl @@ -12,7 +12,6 @@ end include("inplace.jl") include("MatrixAlgebra.jl") include("bituple.jl") -include("BaseExtensions/BaseExtensions.jl") include("permutedimsadd.jl") include("matricize.jl") include("contract/contractalgorithm.jl") diff --git a/src/bituple.jl b/src/bituple.jl index 317227aa..dd04f621 100644 --- a/src/bituple.jl +++ b/src/bituple.jl @@ -1,8 +1,9 @@ -# A two-block tuple: a flat tuple split into a codomain group `t1` and a domain group `t2`. -# This is the only blocked-tuple shape matricization, contraction, and factorizations ever use. -# When its entries are `Int`s forming a permutation it acts as a "biperm"; the permutation -# property is validated at construction by the perm builders (like `permutedims(a, perm)`), -# not encoded in a separate type. +# A two-block tuple: a flat tuple carrying an extra codomain/domain split. It acts like the +# flat tuple `(t1..., t2...)` for iteration, indexing, and `length`, with the split exposed only +# through the `t1` and `t2` fields (the way `Pair` exposes its two halves through fields rather +# than a collection interface). When its entries are `Int`s forming a permutation it acts as a +# "biperm"; whether it is a valid permutation is the concern of the operation using it as one +# (e.g. `matricize`), not of the type. unval(::Val{N}) where {N} = N @@ -11,12 +12,10 @@ struct BiTuple{N1, N2, T1 <: NTuple{N1, Any}, T2 <: NTuple{N2, Any}} t2::T2 end -# Accessors (a small BlockArrays-inspired surface). -blocks(bt::BiTuple) = (bt.t1, bt.t2) -firstblock(bt::BiTuple) = bt.t1 -lastblock(bt::BiTuple) = bt.t2 -blocklengths(::BiTuple{N1, N2}) where {N1, N2} = (N1, N2) -blocklength(::BiTuple) = 2 +# Split a flat tuple into a first block of length `N1` and the remaining second block. +function BiTuple(t::NTuple{N, Any}, ::Val{N1}) where {N, N1} + return BiTuple(ntuple(i -> t[i], Val(N1)), ntuple(i -> t[N1 + i], Val(N - N1))) +end Base.Tuple(bt::BiTuple) = (bt.t1..., bt.t2...) Base.length(::BiTuple{N1, N2}) where {N1, N2} = N1 + N2 @@ -27,56 +26,20 @@ function Base.eltype(::Type{<:BiTuple{<:Any, <:Any, T1, T2}}) where {T1, T2} end function Base.show(io::IO, bt::BiTuple) - return print(io, "tuplemortar(", blocks(bt), ")") -end - -function Base.invperm(bt::BiTuple{N1, N2}) where {N1, N2} - ip = invperm(Tuple(bt)) - return BiTuple(ntuple(i -> ip[i], Val(N1)), ntuple(i -> ip[N1 + i], Val(N2))) + return print(io, "BiTuple(", bt.t1, ", ", bt.t2, ")") end -# -# Constructors -# - -# Axis bituple: split a tuple-of-tuples into the two groups verbatim. -tuplemortar(tt::Tuple{Tuple, Tuple}) = BiTuple(tt[1], tt[2]) +Base.:(==)(a::BiTuple, b::BiTuple) = a.t1 == b.t1 && a.t2 == b.t2 +Base.hash(bt::BiTuple, h::UInt) = hash(bt.t2, hash(bt.t1, hash(:BiTuple, h))) -# Permutation bituples: validate at runtime that the flat tuple is a permutation. -function permmortar(permblocks::Tuple{Tuple{Vararg{Int}}, Tuple{Vararg{Int}}}) - bt = BiTuple(permblocks[1], permblocks[2]) - @assert isperm(Tuple(bt)) - return bt -end - -# Split a flat permutation into a codomain block of length `blocklengths[1]` and the rest. -function blockedperm(perm::Tuple{Vararg{Int}}, blocklengths::Tuple{Int, Int}) - l1 = blocklengths[1] - return permmortar((perm[1:l1], perm[(l1 + 1):end])) -end - -# Two-block vcat builder (the contraction path's entry point). -function blockedpermvcat(block1::Tuple{Vararg{Int}}, block2::Tuple{Vararg{Int}}) - return permmortar((block1, block2)) -end - -# `indexin`-based builder: locate each group's labels within `collection`. -function blockedperm_indexin(collection, sub1, sub2) - return permmortar( - ( - BaseExtensions.indexin(sub1, collection), - BaseExtensions.indexin(sub2, collection), - ) - ) -end - -# Trivial (identity) biperm with the codomain/domain split given by `Val`s. Known to be a -# permutation by construction, so it skips the `isperm` check on this hot path. -function trivialbiperm(::Val{N1}, ::Val{N}) where {N1, N} - return BiTuple(ntuple(identity, Val(N1)), ntuple(i -> N1 + i, Val(N - N1))) -end +Base.invperm(bt::BiTuple{N1}) where {N1} = BiTuple(invperm(Tuple(bt)), Val(N1)) -# Bipartition a collection according to a biperm (out-of-place, blocked). -function blockpermute(v, bt::BiTuple) - return BiTuple(map(i -> v[i], bt.t1), map(i -> v[i], bt.t2)) +# Partition `v` into two groups. The partition is specified either by a split length (take +# the first `length1` entries in order, then the rest), by two index groups `t1`/`t2`, or by +# a `BiTuple` of index groups. +function bipartition(t::Tuple, length1::Val) + bt = BiTuple(t, length1) + return bt.t1, bt.t2 end +bipartition(v, t1::Tuple, t2::Tuple) = (map(i -> v[i], t1), map(i -> v[i], t2)) +bipartition(v, bt::BiTuple) = bipartition(v, bt.t1, bt.t2) diff --git a/src/contract/allocate_output.jl b/src/contract/allocate_output.jl index 256e0905..2f71f6c5 100644 --- a/src/contract/allocate_output.jl +++ b/src/contract/allocate_output.jl @@ -47,14 +47,13 @@ function output_axes( a1::AbstractArray, perm1_codomain, perm1_domain, a2::AbstractArray, perm2_codomain, perm2_domain ) - biperm1 = permmortar((perm1_codomain, perm1_domain)) - biperm2 = permmortar((perm2_codomain, perm2_domain)) - biperm_dest = permmortar((perm_dest_codomain, perm_dest_domain)) - axes_codomain, axes_contracted = blocks(blockpermute(axes(a1), biperm1)) - axes_contracted2, axes_domain = blocks(blockpermute(axes(a2), biperm2)) + axes_codomain, axes_contracted = bipartition(axes(a1), perm1_codomain, perm1_domain) + axes_contracted2, axes_domain = bipartition(axes(a2), perm2_codomain, perm2_domain) @assert length.(axes_contracted) == length.(axes_contracted2) - # default: flatten biperm_out - return genperm((axes_codomain..., axes_domain...), Tuple(biperm_dest)) + # default: flatten the destination permutation + return genperm( + (axes_codomain..., axes_domain...), (perm_dest_codomain..., perm_dest_domain...) + ) end # TODO: Use `ArrayLayouts`-like `MulAdd` object, diff --git a/src/contract/blockedperms.jl b/src/contract/blockedperms.jl index 6615279f..7de00b80 100644 --- a/src/contract/blockedperms.jl +++ b/src/contract/blockedperms.jl @@ -1,15 +1,24 @@ -using .BaseExtensions: BaseExtensions +# `Base.indexin` doesn't accept tuples; return the positions of `x` in `y` as a tuple. +function tuple_indexin(x::Tuple, y::AbstractArray) + return Tuple{Vararg{Any, length(x)}}(Base.indexin(x, y)) +end +tuple_indexin(x::Tuple, y) = tuple_indexin(x, collect(y)) + +# Locate two subgroups `sub1`, `sub2` within `collection`, returning their two index groups. +function biindexin(collection, sub1, sub2) + return tuple_indexin(sub1, collection), tuple_indexin(sub2, collection) +end -# default: if no bipartion is specified, all axes to domain +# Split `perm` into a codomain block of length `blocklength1` and a domain block. function biperm(perm, blocklength1::Integer) return biperm(perm, Val(blocklength1)) end function biperm(perm, ::Val{BlockLength1}) where {BlockLength1} length(perm) < BlockLength1 && throw(ArgumentError("Invalid codomain length")) - return blockedperm(Tuple(perm), (BlockLength1, length(perm) - BlockLength1)) + return BiTuple(Tuple(perm), Val(BlockLength1)) end -length_domain(t::BiTuple) = last(blocklengths(t)) +length_domain(t::BiTuple) = length(t.t2) # Assume all dimensions are in the codomain by default length_domain(t) = 0 @@ -26,20 +35,18 @@ function blockedperms(::typeof(contract), dimnames_dest, dimnames1, dimnames2) contracted = Tuple(intersect(dimnames1, dimnames2)) domain = Tuple(setdiff(dimnames2, dimnames1)) - perm_codomain_dest = BaseExtensions.indexin(codomain, dimnames_dest) - perm_domain_dest = BaseExtensions.indexin(domain, dimnames_dest) + perm_codomain_dest = tuple_indexin(codomain, dimnames_dest) + perm_domain_dest = tuple_indexin(domain, dimnames_dest) invbiperm = (perm_codomain_dest..., perm_domain_dest...) biperm_dest = biperm(invperm(invbiperm), length(codomain)) - perm_codomain1 = BaseExtensions.indexin(codomain, dimnames1) - perm_domain1 = BaseExtensions.indexin(contracted, dimnames1) + perm_codomain1 = tuple_indexin(codomain, dimnames1) + perm_domain1 = tuple_indexin(contracted, dimnames1) - perm_codomain2 = BaseExtensions.indexin(contracted, dimnames2) - perm_domain2 = BaseExtensions.indexin(domain, dimnames2) + perm_codomain2 = tuple_indexin(contracted, dimnames2) + perm_domain2 = tuple_indexin(domain, dimnames2) - permblocks1 = (perm_codomain1, perm_domain1) - biperm1 = blockedpermvcat(permblocks1...) - permblocks2 = (perm_codomain2, perm_domain2) - biperm2 = blockedpermvcat(permblocks2...) + biperm1 = BiTuple(perm_codomain1, perm_domain1) + biperm2 = BiTuple(perm_codomain2, perm_domain2) return biperm_dest, biperm1, biperm2 end diff --git a/src/contract/contract.jl b/src/contract/contract.jl index 35b12904..a9b95d65 100644 --- a/src/contract/contract.jl +++ b/src/contract/contract.jl @@ -11,9 +11,9 @@ function contract( ) biperm_dest, biperm1, biperm2 = blockedperms(contract, labels_dest, labels1, labels2) return contract( - blocks(biperm_dest)..., - a1, blocks(biperm1)..., - a2, blocks(biperm2)...; + biperm_dest.t1, biperm_dest.t2, + a1, biperm1.t1, biperm1.t2, + a2, biperm2.t1, biperm2.t2; kwargs... ) end @@ -26,7 +26,8 @@ function contract( ) Ndest_codomain = Val(length(perm1_codomain)) Ndest = Val(length(perm1_codomain) + length(perm2_domain)) - perm_dest_codomain, perm_dest_domain = blocks(trivialbiperm(Ndest_codomain, Ndest)) + perm_dest_codomain, perm_dest_domain = + bipartition(ntuple(identity, Ndest), Ndest_codomain) return contract( perm_dest_codomain, perm_dest_domain, a1, perm1_codomain, perm1_domain, @@ -117,9 +118,9 @@ function contractopadd!( ) biperm_dest, biperm1, biperm2 = blockedperms(contract, labels_dest, labels1, labels2) return contractopadd!( - a_dest, blocks(biperm_dest)..., - op1, a1, blocks(biperm1)..., - op2, a2, blocks(biperm2)..., + a_dest, biperm_dest.t1, biperm_dest.t2, + op1, a1, biperm1.t1, biperm1.t2, + op2, a2, biperm2.t1, biperm2.t2, α, β; kwargs... ) end @@ -175,7 +176,7 @@ function contract( a2::AbstractArray, biperm2::BiTuple; kwargs... ) - return contract(a1, blocks(biperm1)..., a2, blocks(biperm2)...; kwargs...) + return contract(a1, biperm1.t1, biperm1.t2, a2, biperm2.t1, biperm2.t2; kwargs...) end function contract( biperm_dest::BiTuple, @@ -184,9 +185,9 @@ function contract( kwargs... ) return contract( - blocks(biperm_dest)..., - a1, blocks(biperm1)..., - a2, blocks(biperm2)...; + biperm_dest.t1, biperm_dest.t2, + a1, biperm1.t1, biperm1.t2, + a2, biperm2.t1, biperm2.t2; kwargs... ) end @@ -197,9 +198,9 @@ function contract!( kwargs... ) return contract!( - a_dest, blocks(biperm_dest)..., - a1, blocks(biperm1)..., - a2, blocks(biperm2)...; + a_dest, biperm_dest.t1, biperm_dest.t2, + a1, biperm1.t1, biperm1.t2, + a2, biperm2.t1, biperm2.t2; kwargs... ) end @@ -210,9 +211,9 @@ function contractadd!( α::Number, β::Number; kwargs... ) return contractadd!( - a_dest, blocks(biperm_dest)..., - a1, blocks(biperm1)..., - a2, blocks(biperm2)..., + a_dest, biperm_dest.t1, biperm_dest.t2, + a1, biperm1.t1, biperm1.t2, + a2, biperm2.t1, biperm2.t2, α, β; kwargs... ) end @@ -223,9 +224,9 @@ function contractopadd!( α::Number, β::Number; kwargs... ) return contractopadd!( - a_dest, blocks(biperm_dest)..., - op1, a1, blocks(biperm1)..., - op2, a2, blocks(biperm2)..., + a_dest, biperm_dest.t1, biperm_dest.t2, + op1, a1, biperm1.t1, biperm1.t2, + op2, a2, biperm2.t1, biperm2.t2, α, β; kwargs... ) end diff --git a/src/contract/contract_matricize.jl b/src/contract/contract_matricize.jl index 8cdbc389..ceb7e911 100644 --- a/src/contract/contract_matricize.jl +++ b/src/contract/contract_matricize.jl @@ -6,8 +6,8 @@ function contractopadd!( α::Number, β::Number ) biperm_dest = (biperm_dest_codomain..., biperm_dest_domain...) - invperm_codomain, invperm_domain = - blocks(biperm(invperm(biperm_dest), length(biperm1_codomain))) + invbiperm = biperm(invperm(biperm_dest), length(biperm1_codomain)) + invperm_codomain, invperm_domain = invbiperm.t1, invbiperm.t2 check_input( contract!, a_dest, invperm_codomain, invperm_domain, diff --git a/src/factorizations.jl b/src/factorizations.jl index 3f3a1525..4b844f1a 100644 --- a/src/factorizations.jl +++ b/src/factorizations.jl @@ -11,10 +11,9 @@ for f in ( function $f(style::FusionStyle, A::AbstractArray, ndims_codomain::Val; kwargs...) A_mat = matricize(style, A, ndims_codomain) X, Y = MatrixAlgebraKit.$f(A_mat; kwargs...) - biperm = trivialbiperm(ndims_codomain, Val(ndims(A))) - axes_codomain, axes_domain = blocks(blockpermute(axes(A), biperm)) - axes_X = tuplemortar((axes_codomain, (axes(X, 2),))) - axes_Y = tuplemortar(((axes(Y, 1),), axes_domain)) + axes_codomain, axes_domain = bipartition(axes(A), ndims_codomain) + axes_X = BiTuple(axes_codomain, (axes(X, 2),)) + axes_Y = BiTuple((axes(Y, 1),), axes_domain) return unmatricize(style, X, axes_X), unmatricize(style, Y, axes_Y) end function $f(A::AbstractArray, ndims_codomain::Val; kwargs...) @@ -52,24 +51,24 @@ for f in ( style::FusionStyle, A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs... ) - biperm = - blockedperm_indexin(Tuple.((labels_A, labels_codomain, labels_domain))...) - return $f(style, A, blocks(biperm)...; kwargs...) + perm_codomain, perm_domain = + biindexin(Tuple.((labels_A, labels_codomain, labels_domain))...) + return $f(style, A, perm_codomain, perm_domain; kwargs...) end function $f(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) - biperm = - blockedperm_indexin(Tuple.((labels_A, labels_codomain, labels_domain))...) - return $f(A, blocks(biperm)...; kwargs...) + perm_codomain, perm_domain = + biindexin(Tuple.((labels_A, labels_codomain, labels_domain))...) + return $f(A, perm_codomain, perm_domain; kwargs...) end function $f( style::FusionStyle, A::AbstractArray, biperm::BiTuple; kwargs... ) - return $f(style, A, blocks(biperm)...; kwargs...) + return $f(style, A, biperm.t1, biperm.t2; kwargs...) end function $f(A::AbstractArray, biperm::BiTuple; kwargs...) - return $f(A, blocks(biperm)...; kwargs...) + return $f(A, biperm.t1, biperm.t2; kwargs...) end end end @@ -229,11 +228,10 @@ for f in (:svd_compact, :svd_full, :svd_trunc) function $f(style::FusionStyle, A::AbstractArray, ndims_codomain::Val; kwargs...) A_mat = matricize(style, A, ndims_codomain) U, S, Vᴴ = MatrixAlgebraKit.$f(A_mat; kwargs...) - biperm = trivialbiperm(ndims_codomain, Val(ndims(A))) - axes_codomain, axes_domain = blocks(blockpermute(axes(A), biperm)) - axes_U = tuplemortar((axes_codomain, (axes(U, 2),))) - axes_S = tuplemortar(((axes(S, 1),), (axes(S, 2),))) - axes_Vᴴ = tuplemortar(((axes(Vᴴ, 1),), axes_domain)) + axes_codomain, axes_domain = bipartition(axes(A), ndims_codomain) + axes_U = BiTuple(axes_codomain, (axes(U, 2),)) + axes_S = BiTuple((axes(S, 1),), (axes(S, 2),)) + axes_Vᴴ = BiTuple((axes(Vᴴ, 1),), axes_domain) return unmatricize(style, U, axes_U), unmatricize(style, S, axes_S), unmatricize(style, Vᴴ, axes_Vᴴ) @@ -251,9 +249,8 @@ for f in (:eigh_full, :eig_full, :eigh_trunc, :eig_trunc) function $f(style::FusionStyle, A::AbstractArray, ndims_codomain::Val; kwargs...) A_mat = matricize(style, A, ndims_codomain) D, V = MatrixAlgebraKit.$f(A_mat; kwargs...) - biperm = trivialbiperm(ndims_codomain, Val(ndims(A))) - axes_codomain, = blocks(blockpermute(axes(A), biperm)) - axes_V = tuplemortar((axes_codomain, (axes(V, ndims(V)),))) + axes_codomain = first(bipartition(axes(A), ndims_codomain)) + axes_V = BiTuple(axes_codomain, (axes(V, ndims(V)),)) return D, unmatricize(style, V, axes_V) end function $f(A::AbstractArray, ndims_codomain::Val; kwargs...) @@ -442,9 +439,8 @@ left_null function left_null!!(style::FusionStyle, A::AbstractArray, ndims_codomain::Val; kwargs...) A_mat = matricize(style, A, ndims_codomain) N = MatrixAlgebraKit.left_null!(A_mat; kwargs...) - biperm = trivialbiperm(ndims_codomain, Val(ndims(A))) - axes_codomain = first(blocks(blockpermute(axes(A), biperm))) - axes_N = tuplemortar((axes_codomain, (axes(N, 2),))) + axes_codomain = first(bipartition(axes(A), ndims_codomain)) + axes_N = BiTuple(axes_codomain, (axes(N, 2),)) return unmatricize(style, N, axes_N) end function left_null!!(A::AbstractArray, ndims_codomain::Val; kwargs...) @@ -482,9 +478,8 @@ right_null function right_null!!(style::FusionStyle, A::AbstractArray, ndims_codomain::Val; kwargs...) A_mat = matricize(style, A, ndims_codomain) Nᴴ = MatrixAlgebraKit.right_null!(A_mat; kwargs...) - biperm = trivialbiperm(ndims_codomain, Val(ndims(A))) - axes_domain = last(blocks((blockpermute(axes(A), biperm)))) - axes_Nᴴ = tuplemortar(((axes(Nᴴ, 1),), axes_domain)) + axes_domain = last(bipartition(axes(A), ndims_codomain)) + axes_Nᴴ = BiTuple((axes(Nᴴ, 1),), axes_domain) return unmatricize(style, Nᴴ, axes_Nᴴ) end function right_null!!(A::AbstractArray, ndims_codomain::Val; kwargs...) @@ -541,9 +536,8 @@ function gram_eigh_full!!( ) A_mat = matricize(style, A, ndims_codomain) X = MatrixAlgebra.gram_eigh_full!!(A_mat; kwargs...) - biperm = trivialbiperm(ndims_codomain, Val(ndims(A))) - axes_codomain = first(blocks(blockpermute(axes(A), biperm))) - axes_X = tuplemortar((axes_codomain, (axes(X, 2),))) + axes_codomain = first(bipartition(axes(A), ndims_codomain)) + axes_X = BiTuple(axes_codomain, (axes(X, 2),)) return unmatricize(style, X, axes_X) end function gram_eigh_full!!(A::AbstractArray, ndims_codomain::Val; kwargs...) @@ -605,10 +599,9 @@ function gram_eigh_full_with_pinv!!( ) A_mat = matricize(style, A, ndims_codomain) X, Y = MatrixAlgebra.gram_eigh_full_with_pinv!!(A_mat; kwargs...) - biperm = trivialbiperm(ndims_codomain, Val(ndims(A))) - axes_codomain = first(blocks(blockpermute(axes(A), biperm))) - axes_X = tuplemortar((axes_codomain, (axes(X, 2),))) - axes_Y = tuplemortar(((axes(Y, 1),), conj.(axes_codomain))) + axes_codomain = first(bipartition(axes(A), ndims_codomain)) + axes_X = BiTuple(axes_codomain, (axes(X, 2),)) + axes_Y = BiTuple((axes(Y, 1),), conj.(axes_codomain)) return unmatricize(style, X, axes_X), unmatricize(style, Y, axes_Y) end function gram_eigh_full_with_pinv!!(A::AbstractArray, ndims_codomain::Val; kwargs...) @@ -661,9 +654,7 @@ one function one!!(style::FusionStyle, A::AbstractArray, ndims_codomain::Val; kwargs...) A_mat = matricize(style, A, ndims_codomain) MatrixAlgebraKit.one!(A_mat) - biperm = trivialbiperm(ndims_codomain, Val(ndims(A))) - axes_codomain, axes_domain = blocks(blockpermute(axes(A), biperm)) - return unmatricize(style, A_mat, axes_codomain, axes_domain) + return unmatricize(style, A_mat, bipartition(axes(A), ndims_codomain)...) end function one!!(A::AbstractArray, ndims_codomain::Val; kwargs...) return one!!(FusionStyle(A), A, ndims_codomain; kwargs...) diff --git a/src/matricize.jl b/src/matricize.jl index 949e26a1..34800d5a 100644 --- a/src/matricize.jl +++ b/src/matricize.jl @@ -1,4 +1,3 @@ -using .BaseExtensions: _permutedims, _permutedims! using EllipsisNotation: Ellipsis using LinearAlgebra: Diagonal @@ -130,8 +129,7 @@ end function matricize_axes(style::FusionStyle, a::AbstractArray, ndims_codomain::Val) unval(ndims_codomain) ≤ ndims(a) || throw(ArgumentError("Codomain length exceeds number of dimensions.")) - biperm = trivialbiperm(ndims_codomain, Val(ndims(a))) - return matricize_axes(style, a, blocks(blockpermute(axes(a), biperm))...) + return matricize_axes(style, a, bipartition(axes(a), ndims_codomain)...) end function matricize_axes(a::AbstractArray, ndims_codomain::Val) return matricize_axes(FusionStyle(a), a, ndims_codomain) @@ -158,14 +156,9 @@ function allocate_output(::typeof(permutedimsop), op, src::AbstractArray, perm_c T = Base.promote_op(op, eltype(src)) axes_co = map(i -> axes(src, i), perm_co) axes_do = map(i -> axes(src, i), perm_do) - return similar(src, T, tuplemortar((axes_co, axes_do))) + return similar(src, T, BiTuple(axes_co, axes_do)) end -# Inner version takes a list of sub-permutations, overload this one if needed. -# TODO: Remove _permutedims once support for Julia 1.10 is dropped -# define permutedims with a BlockedPermuation. Default is to flatten it. -# TODO: Deprecate `permuteblockeddims` in favor of `bipermutedims`. -# Keeping it here for backwards compatibility. function bipermutedims(a::AbstractArray, perm1, perm2) return permutedimsop(identity, a, perm1, perm2) end @@ -173,12 +166,12 @@ function bipermutedims!(a_dest::AbstractArray, a_src::AbstractArray, perm1, perm return bipermutedimsopadd!(a_dest, identity, a_src, perm1, perm2, true, false) end function bipermutedims(a::AbstractArray, biperm::BiTuple) - return bipermutedims(a, blocks(biperm)...) + return bipermutedims(a, biperm.t1, biperm.t2) end function bipermutedims!( a_dest::AbstractArray, a_src::AbstractArray, biperm::BiTuple ) - return bipermutedims!(a_dest, a_src, blocks(biperm)...) + return bipermutedims!(a_dest, a_src, biperm.t1, biperm.t2) end # ===================================== matricize ======================================== @@ -252,7 +245,7 @@ end function matricize( style::FusionStyle, a::AbstractArray, biperm_dest::BiTuple ) - return matricize(style, a, blocks(biperm_dest)...) + return matricize(style, a, biperm_dest.t1, biperm_dest.t2) end # ==================================== matricizeop ======================================= @@ -311,7 +304,7 @@ end function unmatricize( style::FusionStyle, m::AbstractMatrix, blocked_axes::BiTuple ) - return unmatricize(style, m, blocks(blocked_axes)...) + return unmatricize(style, m, blocked_axes.t1, blocked_axes.t2) end function unmatricize( @@ -324,11 +317,10 @@ function unmatricize( style::FusionStyle, m::AbstractMatrix, axes_dest, invperm_codomain::Tuple{Vararg{Int}}, invperm_domain::Tuple{Vararg{Int}} ) - invbiperm = permmortar((invperm_codomain, invperm_domain)) + invbiperm = BiTuple(invperm_codomain, invperm_domain) length(axes_dest) == length(invbiperm) || throw(ArgumentError("axes do not match permutation")) - blocked_axes = blockpermute(axes_dest, invbiperm) - a12 = unmatricize(style, m, blocked_axes) + a12 = unmatricize(style, m, bipartition(axes_dest, invbiperm)...) biperm_dest = biperm(invperm(invbiperm), length_codomain(axes_dest)) return bipermutedims(a12, biperm_dest) end @@ -340,7 +332,7 @@ function unmatricize( style::FusionStyle, m::AbstractMatrix, axes_dest, invbiperm::BiTuple ) - return unmatricize(style, m, axes_dest, blocks(invbiperm)...) + return unmatricize(style, m, axes_dest, invbiperm.t1, invbiperm.t2) end function unmatricize!( @@ -353,11 +345,10 @@ function unmatricize!( style::FusionStyle, a_dest::AbstractArray, m::AbstractMatrix, invperm_codomain::Tuple{Vararg{Int}}, invperm_domain::Tuple{Vararg{Int}} ) - invbiperm = permmortar((invperm_codomain, invperm_domain)) + invbiperm = BiTuple(invperm_codomain, invperm_domain) ndims(a_dest) == length(invbiperm) || throw(ArgumentError("destination does not match permutation")) - blocked_axes = blockpermute(axes(a_dest), invbiperm) - a_perm = unmatricize(style, m, blocked_axes) + a_perm = unmatricize(style, m, bipartition(axes(a_dest), invbiperm)...) biperm_dest = biperm(invperm(invbiperm), length_codomain(axes(a_dest))) return bipermutedims!(a_dest, a_perm, biperm_dest) end @@ -371,7 +362,7 @@ function unmatricize!( style::FusionStyle, a_dest::AbstractArray, m::AbstractMatrix, invbiperm::BiTuple ) - return unmatricize!(style, a_dest, m, blocks(invbiperm)...) + return unmatricize!(style, a_dest, m, invbiperm.t1, invbiperm.t2) end function unmatricizeadd!( @@ -405,7 +396,7 @@ function unmatricizeadd!( α::Number, β::Number ) return unmatricizeadd!( - style, a_dest, m, blocks(invbiperm)..., α, β + style, a_dest, m, invbiperm.t1, invbiperm.t2, α, β ) end diff --git a/src/matrixfunctions.jl b/src/matrixfunctions.jl index 4d882fc8..d5382c47 100644 --- a/src/matrixfunctions.jl +++ b/src/matrixfunctions.jl @@ -36,8 +36,7 @@ for f in MATRIX_FUNCTIONS function $f(style::FusionStyle, a::AbstractArray, ndims_codomain::Val; kwargs...) a_mat = matricize(style, a, ndims_codomain) fa_mat = Base.$f(a_mat; kwargs...) - biperm = trivialbiperm(ndims_codomain, Val(ndims(a))) - return unmatricize(style, fa_mat, blockpermute(axes(a), biperm)) + return unmatricize(style, fa_mat, bipartition(axes(a), ndims_codomain)...) end function $f(a::AbstractArray, ndims_codomain::Val; kwargs...) return $f(FusionStyle(a), a, ndims_codomain; kwargs...) @@ -64,27 +63,27 @@ for f in MATRIX_FUNCTIONS style::FusionStyle, a::AbstractArray, labels_a, labels_codomain, labels_domain; kwargs... ) - biperm = - blockedperm_indexin(Tuple.((labels_a, labels_codomain, labels_domain))...) - return $f(style, a, blocks(biperm)...; kwargs...) + perm_codomain, perm_domain = + biindexin(Tuple.((labels_a, labels_codomain, labels_domain))...) + return $f(style, a, perm_codomain, perm_domain; kwargs...) end function $f( a::AbstractArray, labels_a, labels_codomain, labels_domain; kwargs... ) - biperm = - blockedperm_indexin(Tuple.((labels_a, labels_codomain, labels_domain))...) - return $f(a, blocks(biperm)...; kwargs...) + perm_codomain, perm_domain = + biindexin(Tuple.((labels_a, labels_codomain, labels_domain))...) + return $f(a, perm_codomain, perm_domain; kwargs...) end function $f( style::FusionStyle, a::AbstractArray, biperm::BiTuple; kwargs... ) - return $f(style, a, blocks(biperm)...; kwargs...) + return $f(style, a, biperm.t1, biperm.t2; kwargs...) end function $f(a::AbstractArray, biperm::BiTuple; kwargs...) - return $f(a, blocks(biperm)...; kwargs...) + return $f(a, biperm.t1, biperm.t2; kwargs...) end end end diff --git a/test/test_basics.jl b/test/test_basics.jl index 962d85f6..7e0b2653 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -1,9 +1,9 @@ import TensorAlgebra using EllipsisNotation: var".." using StableRNGs: StableRNG -using TensorAlgebra: ContractAlgorithm, bipermutedims, bipermutedims!, blockedpermvcat, - contract, contract!, contractadd!, length_codomain, length_domain, matricize, - tuplemortar, unmatricize, unmatricize! +using TensorAlgebra: BiTuple, ContractAlgorithm, bipermutedims, bipermutedims!, contract, + contract!, contractadd!, length_codomain, length_domain, matricize, unmatricize, + unmatricize! using TensorOperations: TensorOperations using Test: @test, @test_broken, @test_throws, @testset @@ -13,7 +13,7 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @testset "TensorAlgebra" begin @testset "misc" begin t = (1, 2, 3) - bt = tuplemortar(((1, 2), (3,))) + bt = BiTuple((1, 2), (3,)) @test length_codomain(t) == 3 @test length_codomain(bt) == 2 @test length_domain(t) == 0 @@ -27,7 +27,7 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @testset "bipermutedims (eltype=$elt)" for elt in elts a = randn(elt, 2, 3, 4, 5) - a_perm = bipermutedims(a, blockedpermvcat((3, 1), (2, 4))) + a_perm = bipermutedims(a, BiTuple((3, 1), (2, 4))) @test a_perm == permutedims(a, (3, 1, 2, 4)) a = randn(elt, 2, 3, 4, 5) @@ -36,7 +36,7 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) a = randn(elt, 2, 3, 4, 5) a_perm = Array{elt}(undef, (4, 2, 3, 5)) - bipermutedims!(a_perm, a, blockedpermvcat((3, 1), (2, 4))) + bipermutedims!(a_perm, a, BiTuple((3, 1), (2, 4))) @test a_perm == permutedims(a, (3, 1, 2, 4)) a = randn(elt, 2, 3, 4, 5) @@ -47,7 +47,7 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @testset "matricize (eltype=$elt)" for elt in elts a = randn(elt, 2, 3, 4, 5) - a_fused = matricize(a, blockedpermvcat((1, 2), (3, 4))) + a_fused = matricize(a, BiTuple((1, 2), (3, 4))) @test eltype(a_fused) === elt @test a_fused ≈ reshape(a, 6, 20) @@ -76,7 +76,7 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @test_throws MethodError matricize(a, (1, 2), (3,), (4,)) @test_throws MethodError matricize(a, (1, 2, 3, 4)) - @test_throws ArgumentError matricize(a, blockedpermvcat((1, 2), (3,))) + @test_throws ArgumentError matricize(a, BiTuple((1, 2), (3,))) v = ones(elt, 2) a_fused = matricize(v, (1,), ()) @@ -128,7 +128,7 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) axes0 = axes(a0) m = reshape(a0, 6, 20) - a = unmatricize(m, tuplemortar((axes0[1:2], axes0[3:4]))) + a = unmatricize(m, BiTuple(axes0[1:2], axes0[3:4])) @test eltype(a) === elt @test a ≈ a0 @@ -136,18 +136,18 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @test eltype(a) === elt @test a ≈ a0 - a = unmatricize(m, axes0, blockedpermvcat((1, 2), (3, 4))) + a = unmatricize(m, axes0, BiTuple((1, 2), (3, 4))) @test eltype(a) === elt @test a ≈ a0 - bp = blockedpermvcat((4, 2), (1, 3)) - bpinv = blockedpermvcat((3, 2), (4, 1)) + bp = BiTuple((4, 2), (1, 3)) + bpinv = BiTuple((3, 2), (4, 1)) a = unmatricize(m, map(i -> axes0[i], bp), bpinv) @test eltype(a) === elt @test a ≈ permutedims(a0, Tuple(bp)) a = similar(a0) - unmatricize!(a, m, blockedpermvcat((1, 2), (3, 4))) + unmatricize!(a, m, BiTuple((1, 2), (3, 4))) @test a ≈ a0 m1 = matricize(a0, bp) @@ -172,8 +172,8 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @test a isa Array{elt, 0} @test a[] == m[1, 1] - @test_throws ArgumentError unmatricize(m, (), blockedpermvcat((1, 2), (3,))) - @test_throws ArgumentError unmatricize!(m, m, blockedpermvcat((1, 2), (3,))) + @test_throws ArgumentError unmatricize(m, (), BiTuple((1, 2), (3,))) + @test_throws ArgumentError unmatricize!(m, m, BiTuple((1, 2), (3,))) end alg_tensoroperations = ContractAlgorithm(TensorOperations.StridedBLAS()) @@ -231,9 +231,9 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @test a_dest ≈ a_dest_tensoroperations # Specify with bituple - a_dest = contract(tuplemortar((labels_dest, ())), a1, labels1, a2, labels2) + a_dest = contract(BiTuple(labels_dest, ()), a1, labels1, a2, labels2) @test a_dest ≈ a_dest_tensoroperations - a_dest = contract(tuplemortar(((), labels_dest)), a1, labels1, a2, labels2) + a_dest = contract(BiTuple((), labels_dest), a1, labels1, a2, labels2) @test a_dest ≈ a_dest_tensoroperations a_dest = contract(labels_dest′, a1, labels1, a2, labels2) a_dest_tensoroperations = contract( diff --git a/test/test_bituple.jl b/test/test_bituple.jl index 4b842fa6..469a504b 100644 --- a/test/test_bituple.jl +++ b/test/test_bituple.jl @@ -1,59 +1,50 @@ -using TensorAlgebra: BiTuple, blockedperm, blockedperm_indexin, blockedpermvcat, - blocklength, blocklengths, blockpermute, blocks, firstblock, lastblock, permmortar, - trivialbiperm, tuplemortar -using Test: @test, @test_throws, @testset +using TensorAlgebra: BiTuple, bipartition, tuple_indexin +using Test: @test, @testset using TestExtras: @constinferred -@testset "BiTuple (axis bituple)" begin - bt = @constinferred tuplemortar(((true, 'a'), (2.0,))) +@testset "BiTuple" begin + bt = @constinferred BiTuple((true, 'a'), (2.0,)) @test bt isa BiTuple{2, 1} + @test bt.t1 == (true, 'a') + @test bt.t2 == (2.0,) @test (@constinferred Tuple(bt)) == (true, 'a', 2.0) - @test (@constinferred blocks(bt)) == ((true, 'a'), (2.0,)) - @test (@constinferred firstblock(bt)) == (true, 'a') - @test (@constinferred lastblock(bt)) == (2.0,) - @test (@constinferred blocklengths(bt)) == (2, 1) - @test blocklength(bt) == 2 @test length(bt) == 3 + # Acts like the flat tuple for indexing and iteration. @test bt[1] == true @test bt[3] == 2.0 @test collect(bt) == [true, 'a', 2.0] - bt_int = tuplemortar(((1,), (2, 3))) - @test eltype(bt_int) === Int + @test eltype(BiTuple((1,), (2, 3))) === Int # Empty blocks are allowed. - bt0 = @constinferred tuplemortar(((1,), ())) - @test blocks(bt0) == ((1,), ()) - @test blocklengths(bt0) == (1, 0) + bt0 = @constinferred BiTuple((1,), ()) + @test bt0.t1 == (1,) + @test bt0.t2 == () @test Tuple(bt0) == (1,) + + # Split constructor: split a flat tuple at the given codomain length. + @test (@constinferred BiTuple((3, 4, 5, 2, 1), Val(3))) == BiTuple((3, 4, 5), (2, 1)) + + # Equality compares the two blocks. + @test BiTuple((1, 2), (3,)) == BiTuple((1, 2), (3,)) + @test BiTuple((1, 2), (3,)) != BiTuple((1,), (2, 3)) end -@testset "BiTuple (biperm)" begin - p = @constinferred permmortar(((3, 4, 5), (2, 1))) +@testset "biperm" begin + p = BiTuple((3, 4, 5), (2, 1)) @test Tuple(p) === (3, 4, 5, 2, 1) @test isperm(Tuple(p)) - @test blocks(p) == ((3, 4, 5), (2, 1)) - @test blocklengths(p) == (3, 2) - @test p == blockedpermvcat((3, 4, 5), (2, 1)) - @test p == blockedperm((3, 4, 5, 2, 1), (3, 2)) - @test Tuple(@constinferred invperm(p)) == invperm(Tuple(p)) - - # The perm builders validate that the flat tuple is a permutation. - @test_throws AssertionError permmortar(((3, 5), (2, 1))) - @test_throws AssertionError blockedpermvcat((0, 1), (2, 3)) - - # Trivial biperm: identity split into codomain/domain, built type-stably. - tb = @constinferred trivialbiperm(Val(2), Val(4)) - @test blocks(tb) == ((1, 2), (3, 4)) - @test Tuple(tb) == (1, 2, 3, 4) - @test blocks(@constinferred trivialbiperm(Val(0), Val(2))) == ((), (1, 2)) - - # Locate two label groups within a collection. - p = blockedperm_indexin(("a", "b", "c", "d"), ("c", "a"), ("b", "d")) - @test p == blockedpermvcat((3, 1), (2, 4)) - - # blockpermute splits a collection according to a biperm. - bp = blockpermute((10, 20, 30, 40), trivialbiperm(Val(1), Val(4))) - @test bp isa BiTuple - @test blocks(bp) == ((10,), (20, 30, 40)) + @test (@constinferred invperm(p)) isa BiTuple{3, 2} + @test Tuple(invperm(p)) == invperm(Tuple(p)) + + # `bipartition` splits a flat tuple in place (no permutation). + @test (@constinferred bipartition((3, 4, 5, 2, 1), Val(3))) == ((3, 4, 5), (2, 1)) + @test bipartition((10, 20), Val(0)) == ((), (10, 20)) + + # `tuple_indexin` locates labels within a collection. + @test tuple_indexin(("c", "a"), ("a", "b", "c", "d")) == (3, 1) + + # `bipartition` splits a collection by a biperm or by two index groups. + @test bipartition((10, 20, 30, 40), BiTuple((1,), (2, 3, 4))) == ((10,), (20, 30, 40)) + @test bipartition((10, 20, 30, 40), (1,), (2, 3, 4)) == ((10,), (20, 30, 40)) end diff --git a/test/test_mooncakeext.jl b/test/test_mooncakeext.jl index a7c9dcae..08493c00 100644 --- a/test/test_mooncakeext.jl +++ b/test/test_mooncakeext.jl @@ -1,9 +1,8 @@ using Mooncake: Mooncake using Random: Random using TensorAlgebra: BiTuple, ContractAlgorithm, DefaultContractAlgorithm, Matricize, - allocate_output, biperm, blockedperms, blocks, check_input, contract, contract!, - contract_labels, contractadd!, default_contract_algorithm, permmortar, - select_contract_algorithm + allocate_output, biperm, blockedperms, check_input, contract, contract!, + contract_labels, contractadd!, default_contract_algorithm, select_contract_algorithm using Test: @test, @testset @testset "MooncakeExt" begin @@ -22,16 +21,17 @@ using Test: @test, @testset dest = randn(elt, (2, 2)) a1 = randn(elt, (2, 2)) a2 = randn(elt, (2, 2)) - biperm_dest = permmortar(((1,), (2,))) - biperm1 = permmortar(((1,), (2,))) - biperm2 = permmortar(((1,), (2,))) + biperm_dest = BiTuple((1,), (2,)) + biperm1 = BiTuple((1,), (2,)) + biperm2 = BiTuple((1,), (2,)) labels_dest = (:i, :k) labels1 = (:i, :j) labels2 = (:j, :k) Mooncake.TestUtils.test_rule( - rng, allocate_output, contract, blocks(biperm_dest)..., a1, blocks(biperm1)..., - a2, blocks(biperm2)...; mode, is_primitive + rng, allocate_output, contract, biperm_dest.t1, biperm_dest.t2, a1, biperm1.t1, + biperm1.t2, + a2, biperm2.t1, biperm2.t2; mode, is_primitive ) Mooncake.TestUtils.test_rule(rng, biperm, (1, 2, 3), Val(2); mode, is_primitive) Mooncake.TestUtils.test_rule(rng, biperm, (1, 2, 3), 2; mode, is_primitive) @@ -39,12 +39,13 @@ using Test: @test, @testset rng, blockedperms, contract, labels_dest, labels1, labels2; mode, is_primitive ) Mooncake.TestUtils.test_rule( - rng, check_input, contract, a1, blocks(biperm1)..., a2, blocks(biperm2)...; + rng, check_input, contract, a1, biperm1.t1, biperm1.t2, a2, biperm2.t1, + biperm2.t2; mode, is_primitive ) Mooncake.TestUtils.test_rule( - rng, check_input, contract!, dest, blocks(biperm_dest)..., - a1, blocks(biperm1)..., a2, blocks(biperm2)...; mode, is_primitive + rng, check_input, contract!, dest, biperm_dest.t1, biperm_dest.t2, + a1, biperm1.t1, biperm1.t2, a2, biperm2.t1, biperm2.t2; mode, is_primitive ) Mooncake.TestUtils.test_rule( rng, contract_labels, labels1, labels2; mode, is_primitive @@ -67,12 +68,12 @@ using Test: @test, @testset dest = randn(elt, (2, 2)) a1 = randn(elt, (2, 2)) a2 = randn(elt, (2, 2)) - biperm_dest = permmortar(((1,), (2,))) - biperm1 = permmortar(((1,), (2,))) - biperm2 = permmortar(((1,), (2,))) + biperm_dest = BiTuple((1,), (2,)) + biperm1 = BiTuple((1,), (2,)) + biperm2 = BiTuple((1,), (2,)) Mooncake.TestUtils.test_rule( - rng, contractadd!, dest, blocks(biperm_dest)..., - a1, blocks(biperm1)..., a2, blocks(biperm2)..., α, β; + rng, contractadd!, dest, biperm_dest.t1, biperm_dest.t2, + a1, biperm1.t1, biperm1.t2, a2, biperm2.t1, biperm2.t2, α, β; atol, rtol, mode, is_primitive ) end From e2aeab716bbd143db2b07cc481bfd8bad155dd9e Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Wed, 24 Jun 2026 12:34:29 -0400 Subject: [PATCH 3/9] Bump TensorAlgebra compat to 0.11 in the test, docs, and examples projects The subproject manifests still pinned 0.10, which fails resolution and CheckCompatBounds now that the package is 0.11.0. --- docs/Project.toml | 2 +- examples/Project.toml | 2 +- test/Project.toml | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/Project.toml b/docs/Project.toml index e9042548..9b62d571 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -11,4 +11,4 @@ path = ".." Documenter = "1.8.1" ITensorFormatter = "0.2.27" Literate = "2.20.1" -TensorAlgebra = "0.10" +TensorAlgebra = "0.11" diff --git a/examples/Project.toml b/examples/Project.toml index b82572ad..03d310ff 100644 --- a/examples/Project.toml +++ b/examples/Project.toml @@ -5,4 +5,4 @@ TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" path = ".." [compat] -TensorAlgebra = "0.10" +TensorAlgebra = "0.11" diff --git a/test/Project.toml b/test/Project.toml index 35dc1191..7732b69f 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -34,7 +34,7 @@ Random = "1.10" SafeTestsets = "0.1" StableRNGs = "1.0.2" Suppressor = "0.2" -TensorAlgebra = "0.10" +TensorAlgebra = "0.11" TensorOperations = "5.1.4" Test = "1.10" TestExtras = "0.3.1" From 9a8b00c0c03b29d2f7cc492d7d8be6c3ea273858 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Wed, 24 Jun 2026 12:38:39 -0400 Subject: [PATCH 4/9] Rename blockedperms to biperms The function returns the codomain/domain biperms for a contraction, so the name now matches what it produces, in line with the BiTuple vocabulary. Renames the file to match. --- ext/TensorAlgebraMooncakeExt/TensorAlgebraMooncakeExt.jl | 4 ++-- src/TensorAlgebra.jl | 2 +- src/contract/{blockedperms.jl => biperms.jl} | 2 +- src/contract/contract.jl | 4 ++-- test/test_mooncakeext.jl | 6 +++--- 5 files changed, 9 insertions(+), 9 deletions(-) rename src/contract/{blockedperms.jl => biperms.jl} (96%) diff --git a/ext/TensorAlgebraMooncakeExt/TensorAlgebraMooncakeExt.jl b/ext/TensorAlgebraMooncakeExt/TensorAlgebraMooncakeExt.jl index 9db0d909..c1e23c60 100644 --- a/ext/TensorAlgebraMooncakeExt/TensorAlgebraMooncakeExt.jl +++ b/ext/TensorAlgebraMooncakeExt/TensorAlgebraMooncakeExt.jl @@ -1,7 +1,7 @@ module TensorAlgebraMooncakeExt using Mooncake: Mooncake, @zero_derivative, DefaultCtx -using TensorAlgebra: BiTuple, ContractAlgorithm, allocate_output, biperm, blockedperms, +using TensorAlgebra: BiTuple, ContractAlgorithm, allocate_output, biperm, biperms, check_input, contract, contract!, contract_labels, default_contract_algorithm, select_contract_algorithm @@ -12,7 +12,7 @@ Mooncake.tangent_type(::Type{<:ContractAlgorithm}) = Mooncake.NoTangent typeof(allocate_output), typeof(contract), Any, Any, Any, Any, Any, Any, Any, Any, } @zero_derivative DefaultCtx Tuple{typeof(biperm), Any, Any} -@zero_derivative DefaultCtx Tuple{typeof(blockedperms), typeof(contract), Any, Any, Any} +@zero_derivative DefaultCtx Tuple{typeof(biperms), typeof(contract), Any, Any, Any} @zero_derivative DefaultCtx Tuple{ typeof(check_input), typeof(contract), Any, Any, Any, Any, Any, Any, } diff --git a/src/TensorAlgebra.jl b/src/TensorAlgebra.jl index 76fe137c..e82a8e3e 100644 --- a/src/TensorAlgebra.jl +++ b/src/TensorAlgebra.jl @@ -17,7 +17,7 @@ include("matricize.jl") include("contract/contractalgorithm.jl") include("contract/contract.jl") include("contract/contract_labels.jl") -include("contract/blockedperms.jl") +include("contract/biperms.jl") include("contract/allocate_output.jl") include("contract/contract_matricize.jl") include("factorizations.jl") diff --git a/src/contract/blockedperms.jl b/src/contract/biperms.jl similarity index 96% rename from src/contract/blockedperms.jl rename to src/contract/biperms.jl index 7de00b80..c2ea147b 100644 --- a/src/contract/blockedperms.jl +++ b/src/contract/biperms.jl @@ -25,7 +25,7 @@ length_domain(t) = 0 length_codomain(t) = length(t) - length_domain(t) # codomain <-- domain -function blockedperms(::typeof(contract), dimnames_dest, dimnames1, dimnames2) +function biperms(::typeof(contract), dimnames_dest, dimnames1, dimnames2) dimnames = collect(Iterators.flatten((dimnames_dest, dimnames1, dimnames2))) for i in unique(dimnames) count(==(i), dimnames) == 2 || throw(ArgumentError("Invalid contraction labels")) diff --git a/src/contract/contract.jl b/src/contract/contract.jl index a9b95d65..044e4270 100644 --- a/src/contract/contract.jl +++ b/src/contract/contract.jl @@ -9,7 +9,7 @@ end function contract( labels_dest, a1::AbstractArray, labels1, a2::AbstractArray, labels2; kwargs... ) - biperm_dest, biperm1, biperm2 = blockedperms(contract, labels_dest, labels1, labels2) + biperm_dest, biperm1, biperm2 = biperms(contract, labels_dest, labels1, labels2) return contract( biperm_dest.t1, biperm_dest.t2, a1, biperm1.t1, biperm1.t2, @@ -116,7 +116,7 @@ function contractopadd!( α::Number, β::Number; kwargs... ) - biperm_dest, biperm1, biperm2 = blockedperms(contract, labels_dest, labels1, labels2) + biperm_dest, biperm1, biperm2 = biperms(contract, labels_dest, labels1, labels2) return contractopadd!( a_dest, biperm_dest.t1, biperm_dest.t2, op1, a1, biperm1.t1, biperm1.t2, diff --git a/test/test_mooncakeext.jl b/test/test_mooncakeext.jl index 08493c00..8b69ad17 100644 --- a/test/test_mooncakeext.jl +++ b/test/test_mooncakeext.jl @@ -1,8 +1,8 @@ using Mooncake: Mooncake using Random: Random using TensorAlgebra: BiTuple, ContractAlgorithm, DefaultContractAlgorithm, Matricize, - allocate_output, biperm, blockedperms, check_input, contract, contract!, - contract_labels, contractadd!, default_contract_algorithm, select_contract_algorithm + allocate_output, biperm, biperms, check_input, contract, contract!, contract_labels, + contractadd!, default_contract_algorithm, select_contract_algorithm using Test: @test, @testset @testset "MooncakeExt" begin @@ -36,7 +36,7 @@ using Test: @test, @testset Mooncake.TestUtils.test_rule(rng, biperm, (1, 2, 3), Val(2); mode, is_primitive) Mooncake.TestUtils.test_rule(rng, biperm, (1, 2, 3), 2; mode, is_primitive) Mooncake.TestUtils.test_rule( - rng, blockedperms, contract, labels_dest, labels1, labels2; mode, is_primitive + rng, biperms, contract, labels_dest, labels1, labels2; mode, is_primitive ) Mooncake.TestUtils.test_rule( rng, check_input, contract, a1, biperm1.t1, biperm1.t2, a2, biperm2.t1, From 44f7d016bcc884b823caccf93542468d1c48019f Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Wed, 24 Jun 2026 12:42:44 -0400 Subject: [PATCH 5/9] Drop leftover block vocabulary from comments and names Renames the stale BlockPermutation comment, the biperm helper's blocklength1 argument, and the unmatricize biaxes parameter to the BiTuple group vocabulary. --- src/bituple.jl | 4 ++-- src/contract/biperms.jl | 12 ++++++------ src/contract/contract.jl | 2 +- src/matricize.jl | 8 ++++---- 4 files changed, 13 insertions(+), 13 deletions(-) diff --git a/src/bituple.jl b/src/bituple.jl index dd04f621..aa7d88db 100644 --- a/src/bituple.jl +++ b/src/bituple.jl @@ -1,4 +1,4 @@ -# A two-block tuple: a flat tuple carrying an extra codomain/domain split. It acts like the +# A bipartitioned tuple: a flat tuple carrying an extra codomain/domain split. It acts like the # flat tuple `(t1..., t2...)` for iteration, indexing, and `length`, with the split exposed only # through the `t1` and `t2` fields (the way `Pair` exposes its two halves through fields rather # than a collection interface). When its entries are `Int`s forming a permutation it acts as a @@ -12,7 +12,7 @@ struct BiTuple{N1, N2, T1 <: NTuple{N1, Any}, T2 <: NTuple{N2, Any}} t2::T2 end -# Split a flat tuple into a first block of length `N1` and the remaining second block. +# Split a flat tuple into a first group of length `N1` and the remaining second group. function BiTuple(t::NTuple{N, Any}, ::Val{N1}) where {N, N1} return BiTuple(ntuple(i -> t[i], Val(N1)), ntuple(i -> t[N1 + i], Val(N - N1))) end diff --git a/src/contract/biperms.jl b/src/contract/biperms.jl index c2ea147b..4f181f43 100644 --- a/src/contract/biperms.jl +++ b/src/contract/biperms.jl @@ -9,13 +9,13 @@ function biindexin(collection, sub1, sub2) return tuple_indexin(sub1, collection), tuple_indexin(sub2, collection) end -# Split `perm` into a codomain block of length `blocklength1` and a domain block. -function biperm(perm, blocklength1::Integer) - return biperm(perm, Val(blocklength1)) +# Split `perm` into a codomain group of length `length1` and a domain group. +function biperm(perm, length1::Integer) + return biperm(perm, Val(length1)) end -function biperm(perm, ::Val{BlockLength1}) where {BlockLength1} - length(perm) < BlockLength1 && throw(ArgumentError("Invalid codomain length")) - return BiTuple(Tuple(perm), Val(BlockLength1)) +function biperm(perm, ::Val{Length1}) where {Length1} + length(perm) < Length1 && throw(ArgumentError("Invalid codomain length")) + return BiTuple(Tuple(perm), Val(Length1)) end length_domain(t::BiTuple) = length(t.t2) diff --git a/src/contract/contract.jl b/src/contract/contract.jl index 044e4270..7b74b070 100644 --- a/src/contract/contract.jl +++ b/src/contract/contract.jl @@ -170,7 +170,7 @@ function contractopadd!( ) end -# BlockPermutation versions of contract[opadd][!] +# BiTuple versions of contract[opadd][!] function contract( a1::AbstractArray, biperm1::BiTuple, a2::AbstractArray, biperm2::BiTuple; diff --git a/src/matricize.jl b/src/matricize.jl index 34800d5a..7ff5439b 100644 --- a/src/matricize.jl +++ b/src/matricize.jl @@ -298,13 +298,13 @@ function unmatricize( return unmatricize(FusionStyle(m), m, axes_codomain, axes_domain) end -function unmatricize(m::AbstractMatrix, blocked_axes::BiTuple) - return unmatricize(FusionStyle(m), m, blocked_axes) +function unmatricize(m::AbstractMatrix, biaxes::BiTuple) + return unmatricize(FusionStyle(m), m, biaxes) end function unmatricize( - style::FusionStyle, m::AbstractMatrix, blocked_axes::BiTuple + style::FusionStyle, m::AbstractMatrix, biaxes::BiTuple ) - return unmatricize(style, m, blocked_axes.t1, blocked_axes.t2) + return unmatricize(style, m, biaxes.t1, biaxes.t2) end function unmatricize( From 336ba1f17022a1c44385a29ad40bb4a1ada17472 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Wed, 24 Jun 2026 12:57:37 -0400 Subject: [PATCH 6/9] Make biperm the public group-locating helper biperm(t, t1, t2) locates two partitioning groups within t and returns their position tuples (p1, p2), replacing the internal biindexin. Unlike the general indexin it asserts that the groups partition t, checking the lengths add up, so the result is always a valid bipartitioned permutation. The old split-builder biperm(perm, ::Val) folds into the existing BiTuple constructor and bipartition, and biperms returns plain 2-tuples since its results were always split. biperm, bipartition, and BiTuple are now marked public. --- .../TensorAlgebraMooncakeExt.jl | 2 +- src/TensorAlgebra.jl | 6 ++- src/contract/biperms.jl | 45 +++++++++---------- src/contract/contract.jl | 18 ++++---- src/contract/contract_matricize.jl | 4 +- src/factorizations.jl | 4 +- src/matricize.jl | 4 +- src/matrixfunctions.jl | 4 +- test/test_bituple.jl | 9 +++- test/test_matrixfunctions.jl | 4 +- test/test_mooncakeext.jl | 5 ++- 11 files changed, 57 insertions(+), 48 deletions(-) diff --git a/ext/TensorAlgebraMooncakeExt/TensorAlgebraMooncakeExt.jl b/ext/TensorAlgebraMooncakeExt/TensorAlgebraMooncakeExt.jl index c1e23c60..8913e1cc 100644 --- a/ext/TensorAlgebraMooncakeExt/TensorAlgebraMooncakeExt.jl +++ b/ext/TensorAlgebraMooncakeExt/TensorAlgebraMooncakeExt.jl @@ -11,7 +11,7 @@ Mooncake.tangent_type(::Type{<:ContractAlgorithm}) = Mooncake.NoTangent @zero_derivative DefaultCtx Tuple{ typeof(allocate_output), typeof(contract), Any, Any, Any, Any, Any, Any, Any, Any, } -@zero_derivative DefaultCtx Tuple{typeof(biperm), Any, Any} +@zero_derivative DefaultCtx Tuple{typeof(biperm), Any, Any, Any} @zero_derivative DefaultCtx Tuple{typeof(biperms), typeof(contract), Any, Any, Any} @zero_derivative DefaultCtx Tuple{ typeof(check_input), typeof(contract), Any, Any, Any, Any, Any, Any, diff --git a/src/TensorAlgebra.jl b/src/TensorAlgebra.jl index e82a8e3e..0a414237 100644 --- a/src/TensorAlgebra.jl +++ b/src/TensorAlgebra.jl @@ -6,7 +6,11 @@ export contract, contract!, eig_full, eig_trunc, eig_vals, eigh_full, eigh_trunc right_polar, svd_compact, svd_full, svd_trunc, svd_vals if VERSION >= v"1.11.0-DEV.469" - eval(Meta.parse("public contractopadd!, matricizeop, zero!, scale!, permuteddims")) + eval( + Meta.parse( + "public BiTuple, biperm, bipartition, contractopadd!, matricizeop, zero!, scale!, permuteddims" + ) + ) end include("inplace.jl") diff --git a/src/contract/biperms.jl b/src/contract/biperms.jl index 4f181f43..57e3a980 100644 --- a/src/contract/biperms.jl +++ b/src/contract/biperms.jl @@ -4,18 +4,22 @@ function tuple_indexin(x::Tuple, y::AbstractArray) end tuple_indexin(x::Tuple, y) = tuple_indexin(x, collect(y)) -# Locate two subgroups `sub1`, `sub2` within `collection`, returning their two index groups. -function biindexin(collection, sub1, sub2) - return tuple_indexin(sub1, collection), tuple_indexin(sub2, collection) -end - -# Split `perm` into a codomain group of length `length1` and a domain group. -function biperm(perm, length1::Integer) - return biperm(perm, Val(length1)) -end -function biperm(perm, ::Val{Length1}) where {Length1} - length(perm) < Length1 && throw(ArgumentError("Invalid codomain length")) - return BiTuple(Tuple(perm), Val(Length1)) +""" + biperm(t, t1, t2) -> (p1, p2) + +Locate the groups `t1` and `t2` within `t`, returning the positions of `t1` as +`p1` and the positions of `t2` as `p2`. The groups `t1` and `t2` must partition +`t`, so the concatenation `(p1..., p2...)` is a permutation of `eachindex(t)` and +the pair `(p1, p2)` is a bipartitioned permutation (a "biperm") splitting `t` +into a codomain `p1` and a domain `p2`. +""" +function biperm(t, t1, t2) + length(t1) + length(t2) == length(t) || throw( + ArgumentError( + "groups of lengths $(length(t1)) and $(length(t2)) do not partition a collection of length $(length(t))" + ) + ) + return tuple_indexin(t1, t), tuple_indexin(t2, t) end length_domain(t::BiTuple) = length(t.t2) @@ -35,18 +39,11 @@ function biperms(::typeof(contract), dimnames_dest, dimnames1, dimnames2) contracted = Tuple(intersect(dimnames1, dimnames2)) domain = Tuple(setdiff(dimnames2, dimnames1)) - perm_codomain_dest = tuple_indexin(codomain, dimnames_dest) - perm_domain_dest = tuple_indexin(domain, dimnames_dest) - invbiperm = (perm_codomain_dest..., perm_domain_dest...) - biperm_dest = biperm(invperm(invbiperm), length(codomain)) - - perm_codomain1 = tuple_indexin(codomain, dimnames1) - perm_domain1 = tuple_indexin(contracted, dimnames1) - - perm_codomain2 = tuple_indexin(contracted, dimnames2) - perm_domain2 = tuple_indexin(domain, dimnames2) + perm_codomain_dest, perm_domain_dest = biperm(dimnames_dest, codomain, domain) + invperm_dest = invperm((perm_codomain_dest..., perm_domain_dest...)) + biperm_dest = bipartition(invperm_dest, Val(length(codomain))) - biperm1 = BiTuple(perm_codomain1, perm_domain1) - biperm2 = BiTuple(perm_codomain2, perm_domain2) + biperm1 = biperm(dimnames1, codomain, contracted) + biperm2 = biperm(dimnames2, contracted, domain) return biperm_dest, biperm1, biperm2 end diff --git a/src/contract/contract.jl b/src/contract/contract.jl index 7b74b070..d3ab23b2 100644 --- a/src/contract/contract.jl +++ b/src/contract/contract.jl @@ -9,11 +9,12 @@ end function contract( labels_dest, a1::AbstractArray, labels1, a2::AbstractArray, labels2; kwargs... ) - biperm_dest, biperm1, biperm2 = biperms(contract, labels_dest, labels1, labels2) + (perm_dest_codomain, perm_dest_domain), (perm1_codomain, perm1_domain), + (perm2_codomain, perm2_domain) = biperms(contract, labels_dest, labels1, labels2) return contract( - biperm_dest.t1, biperm_dest.t2, - a1, biperm1.t1, biperm1.t2, - a2, biperm2.t1, biperm2.t2; + perm_dest_codomain, perm_dest_domain, + a1, perm1_codomain, perm1_domain, + a2, perm2_codomain, perm2_domain; kwargs... ) end @@ -116,11 +117,12 @@ function contractopadd!( α::Number, β::Number; kwargs... ) - biperm_dest, biperm1, biperm2 = biperms(contract, labels_dest, labels1, labels2) + (perm_dest_codomain, perm_dest_domain), (perm1_codomain, perm1_domain), + (perm2_codomain, perm2_domain) = biperms(contract, labels_dest, labels1, labels2) return contractopadd!( - a_dest, biperm_dest.t1, biperm_dest.t2, - op1, a1, biperm1.t1, biperm1.t2, - op2, a2, biperm2.t1, biperm2.t2, + a_dest, perm_dest_codomain, perm_dest_domain, + op1, a1, perm1_codomain, perm1_domain, + op2, a2, perm2_codomain, perm2_domain, α, β; kwargs... ) end diff --git a/src/contract/contract_matricize.jl b/src/contract/contract_matricize.jl index ceb7e911..f77f8c3b 100644 --- a/src/contract/contract_matricize.jl +++ b/src/contract/contract_matricize.jl @@ -6,8 +6,8 @@ function contractopadd!( α::Number, β::Number ) biperm_dest = (biperm_dest_codomain..., biperm_dest_domain...) - invbiperm = biperm(invperm(biperm_dest), length(biperm1_codomain)) - invperm_codomain, invperm_domain = invbiperm.t1, invbiperm.t2 + invperm_codomain, invperm_domain = + bipartition(invperm(biperm_dest), Val(length(biperm1_codomain))) check_input( contract!, a_dest, invperm_codomain, invperm_domain, diff --git a/src/factorizations.jl b/src/factorizations.jl index 4b844f1a..0a27b640 100644 --- a/src/factorizations.jl +++ b/src/factorizations.jl @@ -52,12 +52,12 @@ for f in ( labels_A, labels_codomain, labels_domain; kwargs... ) perm_codomain, perm_domain = - biindexin(Tuple.((labels_A, labels_codomain, labels_domain))...) + biperm(Tuple.((labels_A, labels_codomain, labels_domain))...) return $f(style, A, perm_codomain, perm_domain; kwargs...) end function $f(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) perm_codomain, perm_domain = - biindexin(Tuple.((labels_A, labels_codomain, labels_domain))...) + biperm(Tuple.((labels_A, labels_codomain, labels_domain))...) return $f(A, perm_codomain, perm_domain; kwargs...) end diff --git a/src/matricize.jl b/src/matricize.jl index 7ff5439b..35d2f236 100644 --- a/src/matricize.jl +++ b/src/matricize.jl @@ -321,7 +321,7 @@ function unmatricize( length(axes_dest) == length(invbiperm) || throw(ArgumentError("axes do not match permutation")) a12 = unmatricize(style, m, bipartition(axes_dest, invbiperm)...) - biperm_dest = biperm(invperm(invbiperm), length_codomain(axes_dest)) + biperm_dest = BiTuple(Tuple(invperm(invbiperm)), Val(length_codomain(axes_dest))) return bipermutedims(a12, biperm_dest) end @@ -349,7 +349,7 @@ function unmatricize!( ndims(a_dest) == length(invbiperm) || throw(ArgumentError("destination does not match permutation")) a_perm = unmatricize(style, m, bipartition(axes(a_dest), invbiperm)...) - biperm_dest = biperm(invperm(invbiperm), length_codomain(axes(a_dest))) + biperm_dest = BiTuple(Tuple(invperm(invbiperm)), Val(length_codomain(axes(a_dest)))) return bipermutedims!(a_dest, a_perm, biperm_dest) end diff --git a/src/matrixfunctions.jl b/src/matrixfunctions.jl index d5382c47..b0f4d101 100644 --- a/src/matrixfunctions.jl +++ b/src/matrixfunctions.jl @@ -64,7 +64,7 @@ for f in MATRIX_FUNCTIONS labels_a, labels_codomain, labels_domain; kwargs... ) perm_codomain, perm_domain = - biindexin(Tuple.((labels_a, labels_codomain, labels_domain))...) + biperm(Tuple.((labels_a, labels_codomain, labels_domain))...) return $f(style, a, perm_codomain, perm_domain; kwargs...) end function $f( @@ -72,7 +72,7 @@ for f in MATRIX_FUNCTIONS labels_a, labels_codomain, labels_domain; kwargs... ) perm_codomain, perm_domain = - biindexin(Tuple.((labels_a, labels_codomain, labels_domain))...) + biperm(Tuple.((labels_a, labels_codomain, labels_domain))...) return $f(a, perm_codomain, perm_domain; kwargs...) end diff --git a/test/test_bituple.jl b/test/test_bituple.jl index 469a504b..9cb91b80 100644 --- a/test/test_bituple.jl +++ b/test/test_bituple.jl @@ -1,5 +1,5 @@ -using TensorAlgebra: BiTuple, bipartition, tuple_indexin -using Test: @test, @testset +using TensorAlgebra: BiTuple, bipartition, biperm, tuple_indexin +using Test: @test, @test_throws, @testset using TestExtras: @constinferred @testset "BiTuple" begin @@ -44,6 +44,11 @@ end # `tuple_indexin` locates labels within a collection. @test tuple_indexin(("c", "a"), ("a", "b", "c", "d")) == (3, 1) + # `biperm` locates two partitioning groups within a collection. + @test biperm(("a", "b", "c", "d"), ("c", "b"), ("d", "a")) == ((3, 2), (4, 1)) + # The groups must partition the collection. + @test_throws ArgumentError biperm(("a", "b", "c", "d"), ("c", "b"), ("a",)) + # `bipartition` splits a collection by a biperm or by two index groups. @test bipartition((10, 20, 30, 40), BiTuple((1,), (2, 3, 4))) == ((10,), (20, 30, 40)) @test bipartition((10, 20, 30, 40), (1,), (2, 3, 4)) == ((10,), (20, 30, 40)) diff --git a/test/test_matrixfunctions.jl b/test/test_matrixfunctions.jl index 54f9bd58..f9544ef3 100644 --- a/test/test_matrixfunctions.jl +++ b/test/test_matrixfunctions.jl @@ -1,5 +1,5 @@ using StableRNGs: StableRNG -using TensorAlgebra: TensorAlgebra, biperm +using TensorAlgebra: TensorAlgebra, BiTuple using Test: @test, @testset @testset "Matrix functions (eltype=$elt)" for elt in (Float32, ComplexF64) @@ -12,7 +12,7 @@ using Test: @test, @testset for fa in ( TensorAlgebra.$f(a, (:a, :b, :c, :d), (:c, :b), (:d, :a)), TensorAlgebra.$f(a, (3, 2), (4, 1)), - TensorAlgebra.$f(a, biperm((3, 2, 4, 1), Val(2))), + TensorAlgebra.$f(a, BiTuple((3, 2, 4, 1), Val(2))), ) local fa′ = reshape($f(reshape(permutedims(a, (3, 2, 4, 1)), (4, 4))), (2, 2, 2, 2)) diff --git a/test/test_mooncakeext.jl b/test/test_mooncakeext.jl index 8b69ad17..840b5f3c 100644 --- a/test/test_mooncakeext.jl +++ b/test/test_mooncakeext.jl @@ -33,8 +33,9 @@ using Test: @test, @testset biperm1.t2, a2, biperm2.t1, biperm2.t2; mode, is_primitive ) - Mooncake.TestUtils.test_rule(rng, biperm, (1, 2, 3), Val(2); mode, is_primitive) - Mooncake.TestUtils.test_rule(rng, biperm, (1, 2, 3), 2; mode, is_primitive) + Mooncake.TestUtils.test_rule( + rng, biperm, (1, 2, 3), (1, 2), (3,); mode, is_primitive + ) Mooncake.TestUtils.test_rule( rng, biperms, contract, labels_dest, labels1, labels2; mode, is_primitive ) From 6c2065174b290c1f5d5847fe7ba6dda47c36846e Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Wed, 24 Jun 2026 13:12:42 -0400 Subject: [PATCH 7/9] Add the new public names to the exports test BiTuple, biperm, and bipartition are now marked public, so they appear in names(TensorAlgebra) on Julia 1.11+. --- test/test_exports.jl | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/test/test_exports.jl b/test/test_exports.jl index 5726270b..b3e4ad71 100644 --- a/test/test_exports.jl +++ b/test/test_exports.jl @@ -31,7 +31,13 @@ using Test: @test, @testset ] # `public` (Julia 1.11+) adds names to `names()`; include them on 1.11+. if VERSION >= v"1.11.0-DEV.469" - append!(exports, [:contractopadd!, :matricizeop, :zero!, :scale!, :permuteddims]) + append!( + exports, + [ + :BiTuple, :biperm, :bipartition, :contractopadd!, :matricizeop, :zero!, + :scale!, :permuteddims, + ] + ) end @test issetequal(names(TensorAlgebra), exports) From 3caeec59dbf0569db13afb809a9dcc0c4a23fa35 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Wed, 24 Jun 2026 14:11:56 -0400 Subject: [PATCH 8/9] Make BiTuple internal and document bipartition No downstream package names BiTuple, and biperm and bipartition return plain tuples, so BiTuple drops out of the public declaration and stays an internal implementation detail. Adds a docstring for the public bipartition and tightens its index-group method to take a Tuple, matching the split-by-length method and every caller. --- src/TensorAlgebra.jl | 2 +- src/bituple.jl | 21 ++++++++++++++++----- test/test_exports.jl | 2 +- 3 files changed, 18 insertions(+), 7 deletions(-) diff --git a/src/TensorAlgebra.jl b/src/TensorAlgebra.jl index 0a414237..966cef31 100644 --- a/src/TensorAlgebra.jl +++ b/src/TensorAlgebra.jl @@ -8,7 +8,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 BiTuple, biperm, bipartition, contractopadd!, matricizeop, zero!, scale!, permuteddims" + "public biperm, bipartition, contractopadd!, matricizeop, zero!, scale!, permuteddims" ) ) end diff --git a/src/bituple.jl b/src/bituple.jl index aa7d88db..6aa1e8f4 100644 --- a/src/bituple.jl +++ b/src/bituple.jl @@ -34,12 +34,23 @@ Base.hash(bt::BiTuple, h::UInt) = hash(bt.t2, hash(bt.t1, hash(:BiTuple, h))) Base.invperm(bt::BiTuple{N1}) where {N1} = BiTuple(invperm(Tuple(bt)), Val(N1)) -# Partition `v` into two groups. The partition is specified either by a split length (take -# the first `length1` entries in order, then the rest), by two index groups `t1`/`t2`, or by -# a `BiTuple` of index groups. +""" + bipartition(t::Tuple, length1::Val) -> (t1, t2) + bipartition(t::Tuple, group1::Tuple, group2::Tuple) -> (p1, p2) + +Split a flat tuple into two groups, returned as a pair of tuples. + +The first form splits `t` in order, taking the first `length1` entries as `t1` +and the remaining entries as `t2`. The second form gathers the entries of `t` at +the two index groups `group1` and `group2`, returning `p1 = t[group1...]` and +`p2 = t[group2...]`. +""" function bipartition(t::Tuple, length1::Val) bt = BiTuple(t, length1) return bt.t1, bt.t2 end -bipartition(v, t1::Tuple, t2::Tuple) = (map(i -> v[i], t1), map(i -> v[i], t2)) -bipartition(v, bt::BiTuple) = bipartition(v, bt.t1, bt.t2) +function bipartition(t::Tuple, group1::Tuple, group2::Tuple) + return map(i -> t[i], group1), map(i -> t[i], group2) +end +# Split `t` by the two groups of a `BiTuple`. +bipartition(t::Tuple, bt::BiTuple) = bipartition(t, bt.t1, bt.t2) diff --git a/test/test_exports.jl b/test/test_exports.jl index b3e4ad71..3fe1a556 100644 --- a/test/test_exports.jl +++ b/test/test_exports.jl @@ -34,7 +34,7 @@ using Test: @test, @testset append!( exports, [ - :BiTuple, :biperm, :bipartition, :contractopadd!, :matricizeop, :zero!, + :biperm, :bipartition, :contractopadd!, :matricizeop, :zero!, :scale!, :permuteddims, ] ) From 5bbd197509df1f4a15f736fbb8fe145420162ce8 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Wed, 24 Jun 2026 16:48:29 -0400 Subject: [PATCH 9/9] Remove the public BiTuple convenience overloads Every contract, factorization, matrix-function, and matricize/unmatricize method that took a bundled BiTuple did nothing but splat its two groups into the method that takes them as separate tuple arguments, so the BiTuple form was a redundant public entry point. Drop those overloads, leaving the labels and two-tuple forms, and rewrite the few internal factorization call sites that built a BiTuple just to feed unmatricize so they pass the two axis tuples directly. BiTuple stays as an internal bundle for the allocation and permutation helpers. Co-Authored-By: Claude Opus 4.8 --- src/contract/contract.jl | 61 ---------------------------------- src/factorizations.jl | 64 +++++++----------------------------- src/matricize.jl | 57 -------------------------------- src/matrixfunctions.jl | 10 ------ test/test_basics.jl | 42 +++++++++-------------- test/test_matrixfunctions.jl | 3 +- 6 files changed, 28 insertions(+), 209 deletions(-) diff --git a/src/contract/contract.jl b/src/contract/contract.jl index d3ab23b2..7141b6d9 100644 --- a/src/contract/contract.jl +++ b/src/contract/contract.jl @@ -171,64 +171,3 @@ function contractopadd!( ) ) end - -# BiTuple versions of contract[opadd][!] -function contract( - a1::AbstractArray, biperm1::BiTuple, - a2::AbstractArray, biperm2::BiTuple; - kwargs... - ) - return contract(a1, biperm1.t1, biperm1.t2, a2, biperm2.t1, biperm2.t2; kwargs...) -end -function contract( - biperm_dest::BiTuple, - a1::AbstractArray, biperm1::BiTuple, - a2::AbstractArray, biperm2::BiTuple; - kwargs... - ) - return contract( - biperm_dest.t1, biperm_dest.t2, - a1, biperm1.t1, biperm1.t2, - a2, biperm2.t1, biperm2.t2; - kwargs... - ) -end -function contract!( - a_dest::AbstractArray, biperm_dest::BiTuple, - a1::AbstractArray, biperm1::BiTuple, - a2::AbstractArray, biperm2::BiTuple; - kwargs... - ) - return contract!( - a_dest, biperm_dest.t1, biperm_dest.t2, - a1, biperm1.t1, biperm1.t2, - a2, biperm2.t1, biperm2.t2; - kwargs... - ) -end -function contractadd!( - a_dest::AbstractArray, biperm_dest::BiTuple, - a1::AbstractArray, biperm1::BiTuple, - a2::AbstractArray, biperm2::BiTuple, - α::Number, β::Number; kwargs... - ) - return contractadd!( - a_dest, biperm_dest.t1, biperm_dest.t2, - a1, biperm1.t1, biperm1.t2, - a2, biperm2.t1, biperm2.t2, - α, β; kwargs... - ) -end -function contractopadd!( - a_dest::AbstractArray, biperm_dest::BiTuple, - op1, a1::AbstractArray, biperm1::BiTuple, - op2, a2::AbstractArray, biperm2::BiTuple, - α::Number, β::Number; kwargs... - ) - return contractopadd!( - a_dest, biperm_dest.t1, biperm_dest.t2, - op1, a1, biperm1.t1, biperm1.t2, - op2, a2, biperm2.t1, biperm2.t2, - α, β; kwargs... - ) -end diff --git a/src/factorizations.jl b/src/factorizations.jl index 0a27b640..3d5ff050 100644 --- a/src/factorizations.jl +++ b/src/factorizations.jl @@ -12,9 +12,8 @@ for f in ( A_mat = matricize(style, A, ndims_codomain) X, Y = MatrixAlgebraKit.$f(A_mat; kwargs...) axes_codomain, axes_domain = bipartition(axes(A), ndims_codomain) - axes_X = BiTuple(axes_codomain, (axes(X, 2),)) - axes_Y = BiTuple((axes(Y, 1),), axes_domain) - return unmatricize(style, X, axes_X), unmatricize(style, Y, axes_Y) + return unmatricize(style, X, axes_codomain, (axes(X, 2),)), + unmatricize(style, Y, (axes(Y, 1),), axes_domain) end function $f(A::AbstractArray, ndims_codomain::Val; kwargs...) return $f(FusionStyle(A), A, ndims_codomain; kwargs...) @@ -60,16 +59,6 @@ for f in ( biperm(Tuple.((labels_A, labels_codomain, labels_domain))...) return $f(A, perm_codomain, perm_domain; kwargs...) end - - function $f( - style::FusionStyle, A::AbstractArray, - biperm::BiTuple; kwargs... - ) - return $f(style, A, biperm.t1, biperm.t2; kwargs...) - end - function $f(A::AbstractArray, biperm::BiTuple; kwargs...) - return $f(A, biperm.t1, biperm.t2; kwargs...) - end end end @@ -77,7 +66,6 @@ end qr_compact(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> Q, R qr_compact(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> Q, R qr_compact(A::AbstractArray, ndims_codomain::Val; kwargs...) -> Q, R - qr_compact(A::AbstractArray, biperm::BiTuple; kwargs...) -> Q, R Compute the compact QR decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions, where `R` is square. The @@ -96,7 +84,6 @@ qr_compact qr_full(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> Q, R qr_full(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> Q, R qr_full(A::AbstractArray, ndims_codomain::Val; kwargs...) -> Q, R - qr_full(A::AbstractArray, biperm::BiTuple; kwargs...) -> Q, R Compute the full QR decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions, where `Q` is unitary. The @@ -115,7 +102,6 @@ qr_full lq_compact(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> L, Q lq_compact(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> L, Q lq_compact(A::AbstractArray, ndims_codomain::Val; kwargs...) -> L, Q - lq_compact(A::AbstractArray, biperm::BiTuple; kwargs...) -> L, Q Compute the compact LQ decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions, where `L` is square. The @@ -134,7 +120,6 @@ lq_compact lq_full(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> L, Q lq_full(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> L, Q lq_full(A::AbstractArray, ndims_codomain::Val; kwargs...) -> L, Q - lq_full(A::AbstractArray, biperm::BiTuple; kwargs...) -> L, Q Compute the full LQ decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions, where `Q` is unitary. The @@ -153,7 +138,6 @@ lq_full left_polar(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> W, P left_polar(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> W, P left_polar(A::AbstractArray, ndims_codomain::Val; kwargs...) -> W, P - left_polar(A::AbstractArray, biperm::BiTuple; kwargs...) -> W, P Compute the left polar decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. These can be specified either via @@ -171,7 +155,6 @@ left_polar right_polar(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> P, W right_polar(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> P, W right_polar(A::AbstractArray, ndims_codomain::Val; kwargs...) -> P, W - right_polar(A::AbstractArray, biperm::BiTuple; kwargs...) -> P, W Compute the right polar decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. These can be specified either via @@ -189,7 +172,6 @@ right_polar left_orth(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> V, C left_orth(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> V, C left_orth(A::AbstractArray, ndims_codomain::Val; kwargs...) -> V, C - left_orth(A::AbstractArray, biperm::BiTuple; kwargs...) -> V, C Compute the left orthogonal decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. These can be specified either via @@ -207,7 +189,6 @@ left_orth right_orth(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> C, V right_orth(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> C, V right_orth(A::AbstractArray, ndims_codomain::Val; kwargs...) -> C, V - right_orth(A::AbstractArray, biperm::BiTuple; kwargs...) -> C, V Compute the right orthogonal decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. These can be specified either via @@ -229,12 +210,9 @@ for f in (:svd_compact, :svd_full, :svd_trunc) A_mat = matricize(style, A, ndims_codomain) U, S, Vᴴ = MatrixAlgebraKit.$f(A_mat; kwargs...) axes_codomain, axes_domain = bipartition(axes(A), ndims_codomain) - axes_U = BiTuple(axes_codomain, (axes(U, 2),)) - axes_S = BiTuple((axes(S, 1),), (axes(S, 2),)) - axes_Vᴴ = BiTuple((axes(Vᴴ, 1),), axes_domain) - return unmatricize(style, U, axes_U), - unmatricize(style, S, axes_S), - unmatricize(style, Vᴴ, axes_Vᴴ) + return unmatricize(style, U, axes_codomain, (axes(U, 2),)), + unmatricize(style, S, (axes(S, 1),), (axes(S, 2),)), + unmatricize(style, Vᴴ, (axes(Vᴴ, 1),), axes_domain) end function $f(A::AbstractArray, ndims_codomain::Val; kwargs...) return $f(FusionStyle(A), A, ndims_codomain; kwargs...) @@ -250,8 +228,7 @@ for f in (:eigh_full, :eig_full, :eigh_trunc, :eig_trunc) A_mat = matricize(style, A, ndims_codomain) D, V = MatrixAlgebraKit.$f(A_mat; kwargs...) axes_codomain = first(bipartition(axes(A), ndims_codomain)) - axes_V = BiTuple(axes_codomain, (axes(V, ndims(V)),)) - return D, unmatricize(style, V, axes_V) + return D, unmatricize(style, V, axes_codomain, (axes(V, ndims(V)),)) end function $f(A::AbstractArray, ndims_codomain::Val; kwargs...) return $f(FusionStyle(A), A, ndims_codomain; kwargs...) @@ -276,7 +253,6 @@ end svd_compact(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> U, S, Vᴴ svd_compact(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> U, S, Vᴴ svd_compact(A::AbstractArray, ndims_codomain::Val; kwargs...) -> U, S, Vᴴ - svd_compact(A::AbstractArray, biperm::BiTuple; kwargs...) -> U, S, Vᴴ Compute the compact (thin) SVD of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions, where `U` and `Vᴴ` are isometric. @@ -290,7 +266,6 @@ svd_compact svd_full(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> U, S, Vᴴ svd_full(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> U, S, Vᴴ svd_full(A::AbstractArray, ndims_codomain::Val; kwargs...) -> U, S, Vᴴ - svd_full(A::AbstractArray, biperm::BiTuple; kwargs...) -> U, S, Vᴴ Compute the full (thick) SVD of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions, where `U` and `Vᴴ` are unitary. @@ -304,7 +279,6 @@ svd_full svd_trunc(A::AbstractArray, labels_A, labels_codomain, labels_domain; trunc, kwargs...) -> U, S, Vᴴ svd_trunc(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; trunc, kwargs...) -> U, S, Vᴴ svd_trunc(A::AbstractArray, ndims_codomain::Val; trunc, kwargs...) -> U, S, Vᴴ - svd_trunc(A::AbstractArray, biperm::BiTuple; trunc, kwargs...) -> U, S, Vᴴ Compute the truncated SVD of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. The partition is specified either via @@ -323,7 +297,6 @@ svd_trunc svd_vals(A::AbstractArray, labels_A, labels_codomain, labels_domain) -> S svd_vals(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}) -> S svd_vals(A::AbstractArray, ndims_codomain::Val) -> S - svd_vals(A::AbstractArray, biperm::BiTuple) -> S Compute the singular values of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. The partition is specified either @@ -337,7 +310,6 @@ svd_vals eigh_full(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> D, V eigh_full(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> D, V eigh_full(A::AbstractArray, ndims_codomain::Val; kwargs...) -> D, V - eigh_full(A::AbstractArray, biperm::BiTuple; kwargs...) -> D, V Compute the eigenvalue decomposition of a generic N-dimensional array interpreted as a Hermitian linear map from the domain to the codomain dimensions. The partition is specified @@ -351,7 +323,6 @@ eigh_full eig_full(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> D, V eig_full(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> D, V eig_full(A::AbstractArray, ndims_codomain::Val; kwargs...) -> D, V - eig_full(A::AbstractArray, biperm::BiTuple; kwargs...) -> D, V Compute the eigenvalue decomposition of a generic N-dimensional array interpreted as a general (non-Hermitian) linear map from the domain to the codomain dimensions. The output @@ -366,7 +337,6 @@ eig_full eigh_trunc(A::AbstractArray, labels_A, labels_codomain, labels_domain; trunc, kwargs...) -> D, V eigh_trunc(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; trunc, kwargs...) -> D, V eigh_trunc(A::AbstractArray, ndims_codomain::Val; trunc, kwargs...) -> D, V - eigh_trunc(A::AbstractArray, biperm::BiTuple; trunc, kwargs...) -> D, V Truncated Hermitian eigenvalue decomposition, like [`eigh_full`](@ref) but keeping only the eigenvalues selected by the `trunc` strategy. @@ -379,7 +349,6 @@ eigh_trunc eig_trunc(A::AbstractArray, labels_A, labels_codomain, labels_domain; trunc, kwargs...) -> D, V eig_trunc(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; trunc, kwargs...) -> D, V eig_trunc(A::AbstractArray, ndims_codomain::Val; trunc, kwargs...) -> D, V - eig_trunc(A::AbstractArray, biperm::BiTuple; trunc, kwargs...) -> D, V Truncated general eigenvalue decomposition, like [`eig_full`](@ref) but keeping only the eigenvalues selected by the `trunc` strategy. @@ -392,7 +361,6 @@ eig_trunc eigh_vals(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> D eigh_vals(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> D eigh_vals(A::AbstractArray, ndims_codomain::Val; kwargs...) -> D - eigh_vals(A::AbstractArray, biperm::BiTuple; kwargs...) -> D Compute the eigenvalues of a generic N-dimensional array interpreted as a Hermitian linear map from the domain to the codomain dimensions. The output is a vector of eigenvalues. @@ -405,7 +373,6 @@ eigh_vals eig_vals(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> D eig_vals(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> D eig_vals(A::AbstractArray, ndims_codomain::Val; kwargs...) -> D - eig_vals(A::AbstractArray, biperm::BiTuple; kwargs...) -> D Compute the eigenvalues of a generic N-dimensional array interpreted as a general (non-Hermitian) linear map from the domain to the codomain dimensions. The output is a @@ -419,7 +386,6 @@ eig_vals left_null(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> N left_null(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> N left_null(A::AbstractArray, ndims_codomain::Val; kwargs...) -> N - left_null(A::AbstractArray, biperm::BiTuple; kwargs...) -> N Compute the left nullspace of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. These can be specified either via @@ -440,8 +406,7 @@ function left_null!!(style::FusionStyle, A::AbstractArray, ndims_codomain::Val; A_mat = matricize(style, A, ndims_codomain) N = MatrixAlgebraKit.left_null!(A_mat; kwargs...) axes_codomain = first(bipartition(axes(A), ndims_codomain)) - axes_N = BiTuple(axes_codomain, (axes(N, 2),)) - return unmatricize(style, N, axes_N) + return unmatricize(style, N, axes_codomain, (axes(N, 2),)) end function left_null!!(A::AbstractArray, ndims_codomain::Val; kwargs...) return left_null!!(FusionStyle(A), A, ndims_codomain; kwargs...) @@ -458,7 +423,6 @@ end right_null(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> Nᴴ right_null(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> Nᴴ right_null(A::AbstractArray, ndims_codomain::Val::Val; kwargs...) -> Nᴴ - right_null(A::AbstractArray, biperm::BiTuple; kwargs...) -> Nᴴ Compute the right nullspace of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. These can be specified either via @@ -479,8 +443,7 @@ function right_null!!(style::FusionStyle, A::AbstractArray, ndims_codomain::Val; A_mat = matricize(style, A, ndims_codomain) Nᴴ = MatrixAlgebraKit.right_null!(A_mat; kwargs...) axes_domain = last(bipartition(axes(A), ndims_codomain)) - axes_Nᴴ = BiTuple((axes(Nᴴ, 1),), axes_domain) - return unmatricize(style, Nᴴ, axes_Nᴴ) + return unmatricize(style, Nᴴ, (axes(Nᴴ, 1),), axes_domain) end function right_null!!(A::AbstractArray, ndims_codomain::Val; kwargs...) return right_null!!(FusionStyle(A), A, ndims_codomain; kwargs...) @@ -497,7 +460,6 @@ end gram_eigh_full(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> X gram_eigh_full(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> X gram_eigh_full(A::AbstractArray, ndims_codomain::Val; kwargs...) -> X - gram_eigh_full(A::AbstractArray, biperm::BiTuple; kwargs...) -> X Gram factorization of a generic N-dimensional array, interpreting it as a Hermitian positive semi-definite linear map from the domain to the codomain @@ -537,8 +499,7 @@ function gram_eigh_full!!( A_mat = matricize(style, A, ndims_codomain) X = MatrixAlgebra.gram_eigh_full!!(A_mat; kwargs...) axes_codomain = first(bipartition(axes(A), ndims_codomain)) - axes_X = BiTuple(axes_codomain, (axes(X, 2),)) - return unmatricize(style, X, axes_X) + return unmatricize(style, X, axes_codomain, (axes(X, 2),)) end function gram_eigh_full!!(A::AbstractArray, ndims_codomain::Val; kwargs...) return gram_eigh_full!!(FusionStyle(A), A, ndims_codomain; kwargs...) @@ -557,7 +518,6 @@ end gram_eigh_full_with_pinv(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> X, Y gram_eigh_full_with_pinv(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> X, Y gram_eigh_full_with_pinv(A::AbstractArray, ndims_codomain::Val; kwargs...) -> X, Y - gram_eigh_full_with_pinv(A::AbstractArray, biperm::BiTuple; kwargs...) -> X, Y Like [`gram_eigh_full`](@ref), but additionally returns `Y ≈ pinv(X)` such that `Y * X ≈ I` on the rank subspace (a left inverse). The codomain axes @@ -600,9 +560,8 @@ function gram_eigh_full_with_pinv!!( A_mat = matricize(style, A, ndims_codomain) X, Y = MatrixAlgebra.gram_eigh_full_with_pinv!!(A_mat; kwargs...) axes_codomain = first(bipartition(axes(A), ndims_codomain)) - axes_X = BiTuple(axes_codomain, (axes(X, 2),)) - axes_Y = BiTuple((axes(Y, 1),), conj.(axes_codomain)) - return unmatricize(style, X, axes_X), unmatricize(style, Y, axes_Y) + return unmatricize(style, X, axes_codomain, (axes(X, 2),)), + unmatricize(style, Y, (axes(Y, 1),), conj.(axes_codomain)) end function gram_eigh_full_with_pinv!!(A::AbstractArray, ndims_codomain::Val; kwargs...) return gram_eigh_full_with_pinv!!(FusionStyle(A), A, ndims_codomain; kwargs...) @@ -621,7 +580,6 @@ end TensorAlgebra.one(A::AbstractArray, labels_A, labels_codomain, labels_domain) -> Id TensorAlgebra.one(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}) -> Id TensorAlgebra.one(A::AbstractArray, ndims_codomain::Val) -> Id - TensorAlgebra.one(A::AbstractArray, biperm::BiTuple) -> Id Construct the identity operator tensor whose shape mirrors `A`, interpreted as a linear map from the domain to the codomain dimensions. The codomain and domain diff --git a/src/matricize.jl b/src/matricize.jl index 35d2f236..6930fc40 100644 --- a/src/matricize.jl +++ b/src/matricize.jl @@ -239,15 +239,6 @@ function matricize( return matricize(style, a, to_permblocks(a, (perm_codomain, perm_domain))...) end -function matricize(a::AbstractArray, biperm_dest::BiTuple) - return matricize(FusionStyle(a), a, biperm_dest) -end -function matricize( - style::FusionStyle, a::AbstractArray, biperm_dest::BiTuple - ) - return matricize(style, a, biperm_dest.t1, biperm_dest.t2) -end - # ==================================== matricizeop ======================================= """ @@ -298,15 +289,6 @@ function unmatricize( return unmatricize(FusionStyle(m), m, axes_codomain, axes_domain) end -function unmatricize(m::AbstractMatrix, biaxes::BiTuple) - return unmatricize(FusionStyle(m), m, biaxes) -end -function unmatricize( - style::FusionStyle, m::AbstractMatrix, biaxes::BiTuple - ) - return unmatricize(style, m, biaxes.t1, biaxes.t2) -end - function unmatricize( m::AbstractMatrix, axes_dest, invperm1::Tuple{Vararg{Int}}, invperm2::Tuple{Vararg{Int}} @@ -325,16 +307,6 @@ function unmatricize( return bipermutedims(a12, biperm_dest) end -function unmatricize(m::AbstractMatrix, axes_dest, invbiperm::BiTuple) - return unmatricize(FusionStyle(m), m, axes_dest, invbiperm) -end -function unmatricize( - style::FusionStyle, m::AbstractMatrix, axes_dest, - invbiperm::BiTuple - ) - return unmatricize(style, m, axes_dest, invbiperm.t1, invbiperm.t2) -end - function unmatricize!( a_dest::AbstractArray, m::AbstractMatrix, invperm_codomain::Tuple{Vararg{Int}}, invperm_domain::Tuple{Vararg{Int}} @@ -353,18 +325,6 @@ function unmatricize!( return bipermutedims!(a_dest, a_perm, biperm_dest) end -function unmatricize!( - a_dest::AbstractArray, m::AbstractMatrix, invbiperm::BiTuple - ) - return unmatricize!(FusionStyle(m), a_dest, m, invbiperm) -end -function unmatricize!( - style::FusionStyle, a_dest::AbstractArray, m::AbstractMatrix, - invbiperm::BiTuple - ) - return unmatricize!(style, a_dest, m, invbiperm.t1, invbiperm.t2) -end - function unmatricizeadd!( a_dest::AbstractArray, m::AbstractMatrix, invperm_codomain::Tuple{Vararg{Int}}, invperm_domain::Tuple{Vararg{Int}}, @@ -383,23 +343,6 @@ function unmatricizeadd!( return add!(a_dest, a, α, β) end -function unmatricizeadd!( - a_dest::AbstractArray, m::AbstractMatrix, - invbiperm::BiTuple, - α::Number, β::Number - ) - return unmatricizeadd!(FusionStyle(a_dest), a_dest, m, invbiperm, α, β) -end -function unmatricizeadd!( - style::FusionStyle, a_dest::AbstractArray, m::AbstractMatrix, - invbiperm::BiTuple, - α::Number, β::Number - ) - return unmatricizeadd!( - style, a_dest, m, invbiperm.t1, invbiperm.t2, α, β - ) -end - # Defaults to ReshapeFusion, a simple reshape struct ReshapeFusion <: FusionStyle end FusionStyle(::Type{<:AbstractArray}) = ReshapeFusion() diff --git a/src/matrixfunctions.jl b/src/matrixfunctions.jl index b0f4d101..7628cc58 100644 --- a/src/matrixfunctions.jl +++ b/src/matrixfunctions.jl @@ -75,15 +75,5 @@ for f in MATRIX_FUNCTIONS biperm(Tuple.((labels_a, labels_codomain, labels_domain))...) return $f(a, perm_codomain, perm_domain; kwargs...) end - - function $f( - style::FusionStyle, a::AbstractArray, - biperm::BiTuple; kwargs... - ) - return $f(style, a, biperm.t1, biperm.t2; kwargs...) - end - function $f(a::AbstractArray, biperm::BiTuple; kwargs...) - return $f(a, biperm.t1, biperm.t2; kwargs...) - end end end diff --git a/test/test_basics.jl b/test/test_basics.jl index 7e0b2653..379a9530 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -47,10 +47,6 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @testset "matricize (eltype=$elt)" for elt in elts a = randn(elt, 2, 3, 4, 5) - a_fused = matricize(a, BiTuple((1, 2), (3, 4))) - @test eltype(a_fused) === elt - @test a_fused ≈ reshape(a, 6, 20) - a_fused = matricize(a, (1, 2), (3, 4)) @test eltype(a_fused) === elt @test a_fused ≈ reshape(a, 6, 20) @@ -76,7 +72,7 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @test_throws MethodError matricize(a, (1, 2), (3,), (4,)) @test_throws MethodError matricize(a, (1, 2, 3, 4)) - @test_throws ArgumentError matricize(a, BiTuple((1, 2), (3,))) + @test_throws ArgumentError matricize(a, (1, 2), (3,)) v = ones(elt, 2) a_fused = matricize(v, (1,), ()) @@ -128,35 +124,34 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) axes0 = axes(a0) m = reshape(a0, 6, 20) - a = unmatricize(m, BiTuple(axes0[1:2], axes0[3:4])) - @test eltype(a) === elt - @test a ≈ a0 - a = unmatricize(m, axes0[1:2], axes0[3:4]) @test eltype(a) === elt @test a ≈ a0 - a = unmatricize(m, axes0, BiTuple((1, 2), (3, 4))) + a = unmatricize(m, axes0, (1, 2), (3, 4)) @test eltype(a) === elt @test a ≈ a0 - bp = BiTuple((4, 2), (1, 3)) - bpinv = BiTuple((3, 2), (4, 1)) - a = unmatricize(m, map(i -> axes0[i], bp), bpinv) + perm_codomain = (4, 2) + perm_domain = (1, 3) + invperm_codomain = (3, 2) + invperm_domain = (4, 1) + perm = (4, 2, 1, 3) + a = unmatricize(m, map(i -> axes0[i], perm), invperm_codomain, invperm_domain) @test eltype(a) === elt - @test a ≈ permutedims(a0, Tuple(bp)) + @test a ≈ permutedims(a0, perm) a = similar(a0) - unmatricize!(a, m, BiTuple((1, 2), (3, 4))) + unmatricize!(a, m, (1, 2), (3, 4)) @test a ≈ a0 - m1 = matricize(a0, bp) - a = unmatricize(m1, axes0, bp) + m1 = matricize(a0, perm_codomain, perm_domain) + a = unmatricize(m1, axes0, perm_codomain, perm_domain) @test a ≈ a0 - a1 = permutedims(a0, Tuple(bp)) + a1 = permutedims(a0, perm) a = similar(a1) - unmatricize!(a, m, bpinv) + unmatricize!(a, m, invperm_codomain, invperm_domain) @test a ≈ a1 a = unmatricize(m, (), axes0) @@ -172,8 +167,8 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @test a isa Array{elt, 0} @test a[] == m[1, 1] - @test_throws ArgumentError unmatricize(m, (), BiTuple((1, 2), (3,))) - @test_throws ArgumentError unmatricize!(m, m, BiTuple((1, 2), (3,))) + @test_throws ArgumentError unmatricize(m, (), (1, 2), (3,)) + @test_throws ArgumentError unmatricize!(m, m, (1, 2), (3,)) end alg_tensoroperations = ContractAlgorithm(TensorOperations.StridedBLAS()) @@ -230,11 +225,6 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) ) @test a_dest ≈ a_dest_tensoroperations - # Specify with bituple - a_dest = contract(BiTuple(labels_dest, ()), a1, labels1, a2, labels2) - @test a_dest ≈ a_dest_tensoroperations - a_dest = contract(BiTuple((), labels_dest), a1, labels1, a2, labels2) - @test a_dest ≈ a_dest_tensoroperations a_dest = contract(labels_dest′, a1, labels1, a2, labels2) a_dest_tensoroperations = contract( labels_dest′, a1, labels1, a2, labels2; alg = alg_tensoroperations diff --git a/test/test_matrixfunctions.jl b/test/test_matrixfunctions.jl index f9544ef3..23ba84ae 100644 --- a/test/test_matrixfunctions.jl +++ b/test/test_matrixfunctions.jl @@ -1,5 +1,5 @@ using StableRNGs: StableRNG -using TensorAlgebra: TensorAlgebra, BiTuple +using TensorAlgebra: TensorAlgebra using Test: @test, @testset @testset "Matrix functions (eltype=$elt)" for elt in (Float32, ComplexF64) @@ -12,7 +12,6 @@ using Test: @test, @testset for fa in ( TensorAlgebra.$f(a, (:a, :b, :c, :d), (:c, :b), (:d, :a)), TensorAlgebra.$f(a, (3, 2), (4, 1)), - TensorAlgebra.$f(a, BiTuple((3, 2, 4, 1), Val(2))), ) local fa′ = reshape($f(reshape(permutedims(a, (3, 2, 4, 1)), (4, 4))), (2, 2, 2, 2))