Alternative blocked permutation interface for contract (#120)

Author: mtfishman

Project.toml

@@ -1,11 +1,9 @@
 name = "TensorAlgebra"
 uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
-version = "0.6.14"
+version = "0.7.0"
 
 [deps]
-ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
-BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
 FunctionImplementations = "7c7cc465-9c6a-495f-bdd1-f42428e86d0c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -15,17 +13,18 @@ TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138"
 
 [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]
-ArrayLayouts = "1.10.4"
 BlockArrays = "1.7.2"
 EllipsisNotation = "1.8"
 FunctionImplementations = "0.3.1, 0.4"

docs/Project.toml

@@ -9,4 +9,4 @@ TensorAlgebra = {path = ".."}
 [compat]
 Documenter = "1.8.1"
 Literate = "2.20.1"
-TensorAlgebra = "0.6"
+TensorAlgebra = "0.7"

examples/Project.toml

@@ -5,4 +5,4 @@ TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 TensorAlgebra = {path = ".."}
 
 [compat]
-TensorAlgebra = "0.6"
+TensorAlgebra = "0.7"

ext/TensorAlgebraBlockArraysExt/TensorAlgebraBlockArraysExt.jl

@@ -0,0 +1,6 @@
+module TensorAlgebraBlockArraysExt
+
+include("blockarrays.jl")
+include("blockedtuple.jl")
+
+end

src/blockarrays.jl …nsorAlgebraBlockArraysExt/blockarrays.jlsrc/blockarrays.jl renamed to ext/TensorAlgebraBlockArraysExt/blockarrays.jl src/blockarrays.jl renamed to ext/TensorAlgebraBlockArraysExt/blockarrays.jl

@@ -1,10 +1,12 @@
-using BlockArrays: AbstractBlockArray, AbstractBlockedUnitRange, BlockedArray, blockedrange,
-    eachblockaxes1, mortar
+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
-FusionStyle(::Type{<:AbstractBlockArray}) = BlockReshapeFusion()
+TensorAlgebra.FusionStyle(::Type{<:AbstractBlockArray}) = BlockReshapeFusion()
 
-function trivial_axis(
+function TensorAlgebra.trivial_axis(
         style::BlockReshapeFusion, side::Val{:codomain}, a::AbstractArray,
         axes_codomain::Tuple{Vararg{AbstractUnitRange}},
         axes_domain::Tuple{Vararg{AbstractUnitRange}},
@@ -16,7 +18,7 @@ function mortar_axis(axs)
         throw(ArgumentError("Only one-based axes are supported"))
     return blockedrange(length.(axs))
 end
-function tensor_product_axis(
+function TensorAlgebra.tensor_product_axis(
         style::BlockReshapeFusion, side::Val{:codomain},
         r1::AbstractUnitRange, r2::AbstractUnitRange,
     )
@@ -26,7 +28,9 @@ function tensor_product_axis(
     blockaxs = vec(map(splat(tensor_product_axis), blockaxpairs))
     return mortar_axis(blockaxs)
 end
-function matricize(style::BlockReshapeFusion, a::AbstractArray, ndims_codomain::Val)
+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
@@ -35,7 +39,7 @@ function matricize(style::BlockReshapeFusion, a::AbstractArray, ndims_codomain::
     return mortar(bs, ax)
 end
 using BlockArrays: blocklengths
-function unmatricize(
+function TensorAlgebra.unmatricize(
         ::BlockReshapeFusion, m::AbstractMatrix,
         axes_codomain::Tuple{Vararg{AbstractUnitRange}},
         axes_domain::Tuple{Vararg{AbstractUnitRange}},
@@ -54,11 +58,11 @@ function unmatricize(
     return mortar(bs, ax)
 end
 
-FusionStyle(::Type{<:BlockedArray}) = ReshapeFusion()
+TensorAlgebra.FusionStyle(::Type{<:BlockedArray}) = ReshapeFusion()
 unblock(a::BlockedArray) = a.blocks
 unblock(a::AbstractBlockArray) = a[Base.OneTo.(size(a))...]
 unblock(a::AbstractArray) = a
-function matricize(::ReshapeFusion, a::BlockedArray, ndims_codomain::Val)
+function TensorAlgebra.matricize(::ReshapeFusion, a::BlockedArray, ndims_codomain::Val)
     return matricize(ReshapeFusion(), unblock(a), ndims_codomain)
 end
 function unmatricize_blocked(
@@ -72,21 +76,21 @@ function unmatricize_blocked(
     )
     return BlockedArray(a, (axes_codomain..., axes_domain...))
 end
-function unmatricize(
+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 unmatricize(
+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 unmatricize(
+function TensorAlgebra.unmatricize(
         style::ReshapeFusion, m::AbstractMatrix,
         axes_codomain::Tuple{Vararg{AbstractBlockedUnitRange}},
         axes_domain::Tuple{AbstractBlockedUnitRange, Vararg{AbstractBlockedUnitRange}},

ext/TensorAlgebraBlockArraysExt/blockedtuple.jl

@@ -0,0 +1,20 @@
+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.BlockedTuple, 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)

ext/TensorAlgebraMooncakeExt/TensorAlgebraMooncakeExt.jl

@@ -9,13 +9,15 @@ Mooncake.tangent_type(::Type{<:AbstractBlockPermutation}) = Mooncake.NoTangent
 Mooncake.tangent_type(::Type{<:ContractAlgorithm}) = Mooncake.NoTangent
 
 @zero_derivative DefaultCtx Tuple{
-    typeof(allocate_output), typeof(contract), Any, Any, Any, Any, Any,
+    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(check_input), typeof(contract), Any, Any, Any, Any}
 @zero_derivative DefaultCtx Tuple{
-    typeof(check_input), typeof(contract!), Any, Any, Any, Any, Any, Any,
+    typeof(check_input), typeof(contract), Any, Any, Any, Any, Any, Any,
+}
+@zero_derivative DefaultCtx Tuple{
+    typeof(check_input), typeof(contract!), Any, Any, Any, Any, Any, Any, Any, Any, Any,
 }
 @zero_derivative DefaultCtx Tuple{typeof(contract_labels), Any, Any}
 @zero_derivative DefaultCtx Tuple{typeof(contract_labels), Any, Any, Any, Any}