[BUG] BlockArraysExtensions breaks TensorAlgebra.contract

[ogauthe]

BlockArraysExtensions breaks TensorAlgebra.contract. All tests introduced in ITensor/TensorAlgebra.jl#34 pass when BlockSparseArrays.jl is not loaded (with the exception of zero-dim output type).

using BlockArrays: BlockArray, blockedrange
using TensorAlgebra: contract

d = blockedrange([2, 3])
a = BlockArray(ones((d,d,d,d)))
a_dest, dimnames_dest = contract(a, (1, -1, 2, -2), a, (2, -3, 1, -4))  # ok

using BlockSparseArrays   # load BlockSparseArraysTensorAlgebraExt and BlockArraysExtension
a_dest, dimnames_dest = contract(a, (1, -1, 2, -2), a, (2, -3, 1, -4))  # now crashes

ERROR: BoundsError: attempt to access Int64 at index [2]
Stacktrace:
  [1] getindex
    @ ./number.jl:98 [inlined]
  [2] #13
    @ ~/.julia/packages/BlockSparseArrays/RsnSQ/src/BlockArraysExtensions/BlockArraysExtensions.jl:281 [inlined]
  [3] map
    @ ./tuple.jl:356 [inlined]
  [4] blockreshape(a::BlockArray{Float64, 4, Array{…}, NTuple{…}}, axes::Tuple{BlockedOneTo{…}, BlockedOneTo{…}})
    @ BlockSparseArrays ~/.julia/packages/BlockSparseArrays/RsnSQ/src/BlockArraysExtensions/BlockArraysExtensions.jl:281
  [5] fusedims
    @ ~/.julia/packages/BlockSparseArrays/RsnSQ/ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl:19 [inlined]
  [6] fusedims
    @ ~/.julia/packages/TensorAlgebra/yUoY0/src/fusedims.jl:33 [inlined]
  [7] fusedims
    @ ~/.julia/packages/TensorAlgebra/yUoY0/src/fusedims.jl:42 [inlined]
  [8] fusedims
    @ ~/.julia/packages/TensorAlgebra/yUoY0/src/fusedims.jl:67 [inlined]
  [9] fusedims
    @ ~/.julia/packages/TensorAlgebra/yUoY0/src/fusedims.jl:72 [inlined]
 [10] contract!(alg::TensorAlgebra.Matricize, a_dest::BlockArray{…}, biperm_dest::TensorAlgebra.BlockedPermutation{…}, a1::BlockArray{…}, biperm1::TensorAlgebra.BlockedPermutation{…}, a2::BlockArray{…}, biperm2::TensorAlgebra.BlockedPermutation{…}, α::Bool, β::Bool)
    @ TensorAlgebra ~/.julia/packages/TensorAlgebra/yUoY0/src/contract/contract_matricize/contract.jl:14
 [11] #contract#33
    @ ~/.julia/packages/TensorAlgebra/yUoY0/src/contract/contract.jl:122 [inlined]
 [12] contract(alg::TensorAlgebra.Matricize, biperm_dest::TensorAlgebra.BlockedPermutation{…}, a1::BlockArray{…}, biperm1::TensorAlgebra.BlockedPermutation{…}, a2::BlockArray{…}, biperm2::TensorAlgebra.BlockedPermutation{…}, α::Bool)
    @ TensorAlgebra ~/.julia/packages/TensorAlgebra/yUoY0/src/contract/contract.jl:111
 [13] contract(alg::TensorAlgebra.Matricize, labels_dest::NTuple{…}, a1::BlockArray{…}, labels1::NTuple{…}, a2::BlockArray{…}, labels2::NTuple{…}, α::Bool; kwargs::@Kwargs{})
    @ TensorAlgebra ~/.julia/packages/TensorAlgebra/yUoY0/src/contract/contract.jl:92
 [14] contract(alg::TensorAlgebra.Matricize, labels_dest::NTuple{…}, a1::BlockArray{…}, labels1::NTuple{…}, a2::BlockArray{…}, labels2::NTuple{…}, α::Bool)
    @ TensorAlgebra ~/.julia/packages/TensorAlgebra/yUoY0/src/contract/contract.jl:81
 [15] contract(alg::TensorAlgebra.Matricize, a1::BlockArray{…}, labels1::NTuple{…}, a2::BlockArray{…}, labels2::NTuple{…}, α::Bool; kwargs::@Kwargs{})
    @ TensorAlgebra ~/.julia/packages/TensorAlgebra/yUoY0/src/contract/contract.jl:49
 [16] contract
    @ ~/.julia/packages/TensorAlgebra/yUoY0/src/contract/contract.jl:39 [inlined]
 [17] #contract#27
    @ ~/.julia/packages/TensorAlgebra/yUoY0/src/contract/contract.jl:36 [inlined]
 [18] contract
    @ ~/.julia/packages/TensorAlgebra/yUoY0/src/contract/contract.jl:27 [inlined]
 [19] contract(a1::BlockArray{…}, labels1::NTuple{…}, a2::BlockArray{…}, labels2::NTuple{…})
    @ TensorAlgebra ~/.julia/packages/TensorAlgebra/yUoY0/src/contract/contract.jl:27
 [20] top-level scope
    @ REPL[7]:1
Some type information was truncated. Use `show(err)` to see complete types.

I find it surprising that TensorAlgebra.contract(::BlockArray) behavior changes whether or not BlockSparseArrays is loaded. I believe we should refactor package extensions to avoid this, probably creating a new TensorAlgebra/BlockArraysExtension. We may then make TensorAlgebra a strong dependency of BlockSparseArray, then the extension will always be loaded to be used with a BlockSparseArray.

[ogauthe]

Investigating. The FusionStyle changes when BlockSparseArrays is loaded.

using BlockArrays: BlockArray, blockedrange
using TensorAlgebra: fusedims, FusionStyle
d = blockedrange([2, 3])
a1 = BlockArray(ones((d,d,d,d)))
FusionStyle(a1)   # TensorAlgebra.ReshapeFusion()
fusedims(a1, (1,2), (3,4)) isa Base.ReshapedArray{Float64, 2, <:BlockedArray{Float64,4}}

using BlockSparseArrays: BlockSparseArray
FusionStyle(a1)    # TensorAlgebra.BlockReshapeFusion()
fusedims(a1, (1,2), (3,4))  # BoundsError

ERROR: BoundsError: attempt to access Int64 at index [2]
Stacktrace:
  [1] getindex
    @ ./number.jl:98 [inlined]
  [2] #13
    @ ~/.julia/packages/BlockSparseArrays/RsnSQ/src/BlockArraysExtensions/BlockArraysExtensions.jl:281 [inlined]
  [3] map
    @ ./tuple.jl:356 [inlined]
  [4] blockreshape(a::BlockArray{Float64, 4, Array{…}, NTuple{…}}, axes::Tuple{BlockedOneTo{…}, BlockedOneTo{…}})
    @ BlockSparseArrays ~/.julia/packages/BlockSparseArrays/RsnSQ/src/BlockArraysExtensions/BlockArraysExtensions.jl:281
  [5] fusedims
    @ ~/.julia/packages/BlockSparseArrays/RsnSQ/ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl:19 [inlined]
  [6] fusedims
    @ ~/Documents/itensor/TensorAlgebra.jl/src/fusedims.jl:34 [inlined]
  [7] fusedims
    @ ~/Documents/itensor/TensorAlgebra.jl/src/fusedims.jl:43 [inlined]
  [8] fusedims
    @ ~/Documents/itensor/TensorAlgebra.jl/src/fusedims.jl:68 [inlined]
  [9] fusedims(a::BlockArray{…}, blockedperm::TensorAlgebra.BlockedPermutation{…})
    @ TensorAlgebra ~/Documents/itensor/TensorAlgebra.jl/src/fusedims.jl:73
 [10] fusedims(::BlockArray{Float64, 4, Array{…}, NTuple{…}}, ::Tuple{Int64, Int64}, ::Tuple{Int64, Int64})
    @ TensorAlgebra ~/Documents/itensor/TensorAlgebra.jl/src/fusedims.jl:51
 [11] top-level scope
    @ REPL[21]:1
Some type information was truncated. Use `show(err)` to see complete types.

I see 2 possibilites: preserving ReshapeFusion() as a fusion style for BlockArray and BlockedArray OR fixing BlockReshapeFusion and maybe moving it to TensorAlgebra.

[mtfishman]

I think BlockedArray should be ReshapeFusion() and BlockArray should be BlockReshapeFusion() (which we should fix). I'm not sure about moving BlockReshapeFusion to TensorAlgebra.jl, I don't want TensorAlgebra.jl depending on BlockArrays.jl. If we could design it as a package extension that might be ok but I don't see how that would work since it requires defining new functions and types that we also want to using in BlockSparseArrays.jl.

[ogauthe]

The BlockReshapeFusion branch calls eachstoredindex, which is a sparse function from SparseArrayBase. I am not sure it is suited for BlockArray. It returns a Vector{CartesianIndex} for a BlockSparseArray, a CartesianIndices for a BlockedArray and a Base.OneTo for a BlockArray. There error comes from nested indexing of a Base.OneTo.

Edit: for non-sparse arrays, eachstoredindex defaults to eachindex

[mtfishman]

Maybe we should have BlockReshapeFusion and BlockSparseReshapeFusion.

Then, BlockSparseReshapeFusion can be implemented in BlockSparseArrays.jl and BlockReshapeFusion can be implemented in a TensorAlgebraBlockArraysExt package extension of TensorAlgebra.jl.

[ogauthe]

Here is how I see things:

# ===  TensorAlgebraGradedUnitRangesExt, extension of TensorAlgebra.jl  ===
function TensorAlgebra.:⊗(a1::AbstractBlockedUnitRange, a2::AbstractBlockedUnitRange)
  return tensor_product(a1, a2)
end

# ===  TensorAlgebraBlockArraysExt, extension of TensorAlgebra.jl  ===
FusionStyle(::AbstractBlockedUnitRange) = BlockReshapeFusion()

fusedims(a::BlockReshapeFusion, a::AbstractArray, axes::AbstractUnitRange...) = blockreshape(a, axes)
blockreshape(a::AbstractBlockArray, axes::AbstractUnitRange) = ...

# ===  BlockSparseArraysTensorAlgebraExt, extension of BlockSparseArrays.jl ===
TensorAlgebra.blockreshape(a::AbstractBlockSparseArray, axes::AbstractUnitRange) = ...

Then we would not need a BlockSparseReshapeFusion style. Introducing it requires a deeper refactor as currently the FusionStyle is inferred from the axes, which may be BlockedOneTo for both BlockSparseArray and BlockArray.

I am unsure about ⊗(a1::AbstractBlockedUnitRange, a2::AbstractBlockedUnitRange) in TensorAlgebraGradedUnitRangesExt. It implies axes(fusedims(::BlockArray) depends on whether or not GradedUnitRanges is loaded. This would then affect the internals (and probably the performances) of contract(::BlockArray).

# not loading GradedUnitRanges
d = blockedrange([2, 3])
a = BlockedArray(ones((d,d,d,d)))
axes(fusedims(a, (d⊗d, d⊗d)))   # (Base.OneTo(25), Base.OneTo(25))

using GradedUnitRanges: GradedUnitRanges
axes(fusedims(a, (d⊗d, d⊗d)))   # (BlockedOneTo([4, 10, 16, 25]), BlockedOneTo([4, 10, 16, 25]))

[mtfishman]

The thing I don't like about that plan is that it would be weird for blockreshape to be defined in TensorAlgebra.jl. Basing the FusionStyle on the axes may have just been a misdesign and probably this issue points to the fact that it should be defined on the array. Or more specifically, it should be defined such that it can be defined on the array, but if it is not defined for a certain array type it can fall back to being defined on the axes.

Regarding ⊗(a1::AbstractBlockedUnitRange, a2::AbstractBlockedUnitRange) being defined in TensorAlgebraGradedUnitRangesExt, I agree that probably isn't right, what about defining TensorAlgebraBlockArraysExt and moving it there instead?

[ogauthe]

Here is a new proposal:

# ===  TensorAlgebra.jl  ===
TensorAlgebra.FusionStyle(a::AbstractBlockSparseArray) = FusionStyle(a, axes(a))
FusionStyle(::AbstractArray, ::Tuple{Vararg{AbstractUnitRange}}) = ReshapeFusion()  # default
# no more combine_fusion_styles

fusedims(a::ReshapeFusion, a::AbstractArray, axes::AbstractUnitRange...) = ...  # default

⊗(a1::AbstractUnitRange, a2::AbstractUnitRange) = sectorless_tensor_product(a1, a2)
sectorless_tensor_product(a1::AbstractUnitRange, a2::AbstractUnitRange)  = Base.OneTo(prod(length.((a1,a2))  # default

# ===  TensorAlgebraGradedUnitRangesExt, extension of TensorAlgebra.jl  ===
TensorAlgebra.:⊗(a1::AbstractGradedUnitRange, a2::AbstractGradedUnitRange) = GradedUnitRanges.tensor_product(a1, a2)

# ===  TensorAlgebraBlockArraysExt, extension of TensorAlgebra.jl  ===
sectorless_tensor_product(a1::AbstractBlockedUnitRange, a2::AbstractBlockedUnitRange)  = ...

# maybe remove and use default ReshapeFusion
FusionStyle(::AbstractBlockArray, ::Tuple{Vararg{AbstractUnitRange}}) = BlockReshapeFusion()
fusedims(::BlockReshapeFusion, a::AbstractArray, axes::AbstractUnitRange...) = ...

# ===  BlockSparseArraysTensorAlgebraExt, extension of BlockSparseArrays.jl ===
TensorAlgebra.FusionStyle(a::AbstractBlockSparseArray, ::Tuple{Vararg{AbstractUnitRange}}) = SparseBlockReshapeFusion()
TensorAlgebra.FusionStyle(a::AbstractBlockSparseArray, ::Tuple{Vararg{AbstractGradedUnitRange}}) = MergeFusion()

fusedims(::SparseBlockReshapeFusion, a::AbstractArray, axes::AbstractUnitRange...) = ...
fusedims(::MergeFusion, a::AbstractArray, axes::AbstractUnitRange...) = ...

# ===  GradedUnitRanges.jl  ===
using TensorAlgebra: sectorless_tensor_product  # now a dependency of GradedUnitRanges
tensor_product(a1::AbstractUnitRange, a2::AbstractUnitRange) = sectorless_tensor_product(a1, a2)
tensor_product(a1::AbstractGradedUnitRange, a2::AbstractGradedUnitRange) = ...

EDIT: we can even remove FusionStyle(::AbstractBlockArray, ::Tuple{Vararg{AbstractUnitRange}}) = BlockReshapeFusion() as the default ReshapeFusion is working. It is just an optimization for BlockArray (not BlockedArray) that we may implement later.

I don't really like how TensorAlgebra becomes a dependency of GradedUnitRanges. There may also be some conflicts on the definition of ⊗. Unfortunately as of now I do not see how to avoid dependency without either code duplication (for the tensor product of non-labelled BlockedOneTo) or unexpected changes in contract(::BlockedArray.

More generally, I think we need to rethink package dependencies and get a clear pictures on who depends on whom, who has tests dependencies on whom, who has package dependencies on whom.

[mtfishman]

Maybe this warrants defining an interface package (say called TensorProducts.jl) where we define ⊗/tensor_product. TensorProducts.jl can have definitions for Base ranges, and then GradedUnitRanges.jl and TensorAlgebra.jl can overload or use them as needed. I.e. GradedUnitRanges.jl can overload ⊗/tensor_product for AbstractGradedUnitRange, and then TensorAlgebra.jl just has to depend on TensorProducts.jl if it wants to use them. Additionally, overloads of ⊗/tensor_product for AbstractBlockedUnitRange can be defined in a TensorProductsBlockArraysExt package extension of TensorProducts.jl.

[ogauthe]

Indeed creating a new package TensorProducts.jl seems the best option. It would also solve all naming conflict for ⊗.

# ===  TensorProducts.jl  ===
⊗(args...) = tensor_product(args...)
tensor_product(a1, a2, as...) = tensor_product(a1, tensor_product(a2, as...))   # no type constraint to accept AbstractSector later

tensor_product(a1::AbstractUnitRange, a2::AbstractUnitRange)  = Base.OneTo(prod(length.((a1,a2))  # default

# ===  TensorProductBlockArraysExt, extension of TensorProducts.jl  ===
tensor_product(a1::AbstractBlockedUnitRange, a2::AbstractBlockedUnitRange)  = ...

# ===  GradedUnitRanges.jl  ===
using TensorProducts: TensorProducts  # now a dependency of GradedUnitRanges
TensorProducts.tensor_product(a1::AbstractGradedUnitRange, a2::AbstractGradedUnitRange) = ...

# ===  SymmetrySectors.jl  ===
using TensorProducts: TensorProducts  # now a dependency of SymmetrySectors
TensorProducts.tensor_product(s1::AbstractSector, s2::AbstractSector) = ...

# ===  TensorAlgebra.jl  ===
FusionStyle(a::AbstractArray) = FusionStyle(a, axes(a))
FusionStyle(::AbstractArray, ::Tuple{Vararg{AbstractUnitRange}}) = ReshapeFusion()  # default
# no more combine_fusion_styles

fusedims(a::ReshapeFusion, a::AbstractArray, axes::AbstractUnitRange...) = ...  # default
# remove unneeded TensorAlgebraGradedUnitRangesExt
# later, we may define TensorAlgebraBlockArraysExt to optimize fusedims(::BlockArray)

# ===  BlockSparseArraysTensorAlgebraExt, extension of BlockSparseArrays.jl ===
TensorAlgebra.FusionStyle(a::AbstractBlockSparseArray, ::Tuple{Vararg{AbstractUnitRange}}) = SparseBlockReshapeFusion()

TensorAlgebra.fusedims(::SparseBlockReshapeFusion, a::AbstractArray, axes::AbstractUnitRange...) = ...

# ===  BlockSparseArraysGradedUnitRanges extension of BlockSparseArrays.jl ===
TensorAlgebra.FusionStyle(a::AbstractBlockSparseArray, ::Tuple{Vararg{AbstractGradedUnitRange}}) = MergeFusion()

TensorAlgebra.fusedims(::MergeFusion, a::AbstractArray, axes::AbstractUnitRange...) = ...

We may even add OneToOne definition in TensorProducts.jl, define tensor_product() = OneToOne() and use it as a singleton dimension to solve the issues mentioned in ITensor/TensorAlgebra.jl#33 and #59.

[mtfishman]

That plans looks good to me. Defining OneToOne in TensorProducts.jl also makes sense to me.

What do you think of renaming BlockSparseFusion to SparseBlockReshapeFusion and SectorMergeFusion to MergeFusion? I see SparseBlockReshapeFusion is based on the naming scheme BlockReshapeFusion but maybe that could just be BlockFusion.

I'd be a bit worried about getting rid of combine_fusion_styles in case an array has mixed axes types (say if a singleton dimension of a symmetric array is represented as OneToOne as we've discussed elsewhere). Also, keeping something like that around could make it easier to handle cases where axes are wrapped, which occurs in NamedDimsArrays.jl since the axes are wrapped as named axes. Having said that, we would want to make sure it doesn't get too complicated, since now we have to account for combining the styles of the array and also the axes. Maybe what we could do is:

FusionStyle(a::AbstractArray) = FusionStyle(a, axes(a))
FusionStyle(a::AbstractArray, ax::Tuple{Vararg{AbstractUnitRange}}) = FusionStyle(a, combine_fusion_styles(FusionStyle.(ax)...))

FusionStyle(ax::AbstractUnitRange) = ReshapeFusion()  # default
FusionStyle(a::AbstractArray, ax_style::FusionStyle) = ReshapeFusion()  # default

so then you can decide the fusion style on the array type and the combined fusion styles of the axes.