[ENHANCEMENT] TensorAlgebra functionalities

[ogauthe]

This issue lists functionalities and feature requests for TensorAlgebra.

Issues

• perf: when β is given to contract, the input array C is permuted twice: first to fit A * B, then back to its initial format. This is suboptimal, especially for non-abelian tensors where permutedims is the bottleneck and contractions are fast.
• contracting zero axis crashes. I think it should return the outer product of two tensors. (Crossing this off since it is now being tracked in [BUG] Outer product tensor contractions are broken #13.)

Feature requests

• contract should allow a biperm as input
• ideally, the compiler should be able to predict contract output type

[ogauthe]

Copied from ITensor/ITensors.jl#1473

I investigated TensorAlgebra.contract output type. Currently the compiler cannot deduce anything and returns Any.
The issue comes from TensorAlgebra.blockedperms at ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/blockedperms.jl:11.

julia> @code_warntype TensorAlgebra.blockedperms(TensorAlgebra.contract, (1,2,4), (1,2,3), (4,5))
MethodInstance for NDTensors.TensorAlgebra.blockedperms(::typeof(NDTensors.TensorAlgebra.contract), ::Tuple{Int64, Int64, Int64}, ::Tuple{Int64, Int64, Int64}, ::Tuple{Int64, Int64})
  from blockedperms(::typeof(NDTensors.TensorAlgebra.contract), dimnames_dest, dimnames1, dimnames2) @ NDTensors.TensorAlgebra ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/blockedperms.jl:11
Arguments
  #self#::Core.Const(NDTensors.TensorAlgebra.blockedperms)
  _::Core.Const(NDTensors.TensorAlgebra.contract)
  dimnames_dest::Tuple{Int64, Int64, Int64}
  dimnames1::Tuple{Int64, Int64, Int64}
  dimnames2::Tuple{Int64, Int64}
Locals
  biperm2::Any
  permblocks2::Tuple{Any, Any}
  biperm1::Any
  permblocks1::Tuple{Any, Any}
  biperm_dest::Any
  permblocks_dest::Tuple{Any, Any}
  perm_domain2::Any
  perm_codomain2::Any
  perm_domain1::Any
  perm_codomain1::Any
  perm_domain_dest::Any
  perm_codomain_dest::Any
  domain::Tuple{Vararg{Int64}}
  contracted::Tuple{Vararg{Int64}}
  codomain::Tuple{Vararg{Int64}}
Body::Tuple{Any, Any, Any}
1 ─ %1  = NDTensors.TensorAlgebra.setdiff(dimnames1, dimnames2)::Vector{Int64}
│         (codomain = NDTensors.TensorAlgebra.Tuple(%1))
│   %3  = NDTensors.TensorAlgebra.intersect(dimnames1, dimnames2)::Vector{Int64}
│         (contracted = NDTensors.TensorAlgebra.Tuple(%3))
│   %5  = NDTensors.TensorAlgebra.setdiff(dimnames2, dimnames1)::Vector{Int64}
│         (domain = NDTensors.TensorAlgebra.Tuple(%5))
│   %7  = NDTensors.TensorAlgebra.BaseExtensions.indexin::Core.Const(NDTensors.TensorAlgebra.BaseExtensions.indexin)
│   %8  = codomain::Tuple{Vararg{Int64}}
│         (perm_codomain_dest = (%7)(%8, dimnames_dest))
│   %10 = NDTensors.TensorAlgebra.BaseExtensions.indexin::Core.Const(NDTensors.TensorAlgebra.BaseExtensions.indexin)
│   %11 = domain::Tuple{Vararg{Int64}}
│         (perm_domain_dest = (%10)(%11, dimnames_dest))
│   %13 = NDTensors.TensorAlgebra.BaseExtensions.indexin::Core.Const(NDTensors.TensorAlgebra.BaseExtensions.indexin)
│   %14 = codomain::Tuple{Vararg{Int64}}
│         (perm_codomain1 = (%13)(%14, dimnames1))
│   %16 = NDTensors.TensorAlgebra.BaseExtensions.indexin::Core.Const(NDTensors.TensorAlgebra.BaseExtensions.indexin)
│   %17 = contracted::Tuple{Vararg{Int64}}
│         (perm_domain1 = (%16)(%17, dimnames1))
│   %19 = NDTensors.TensorAlgebra.BaseExtensions.indexin::Core.Const(NDTensors.TensorAlgebra.BaseExtensions.indexin)
│   %20 = contracted::Tuple{Vararg{Int64}}
│         (perm_codomain2 = (%19)(%20, dimnames2))
│   %22 = NDTensors.TensorAlgebra.BaseExtensions.indexin::Core.Const(NDTensors.TensorAlgebra.BaseExtensions.indexin)
│   %23 = domain::Tuple{Vararg{Int64}}
│         (perm_domain2 = (%22)(%23, dimnames2))
│         (permblocks_dest = Core.tuple(perm_codomain_dest, perm_domain_dest))
│   %26 = NDTensors.TensorAlgebra.blockedperm::Core.Const(NDTensors.TensorAlgebra.blockedperm)
│   %27 = !NDTensors.TensorAlgebra.isempty::Core.Const(!isempty)
│   %28 = NDTensors.TensorAlgebra.filter(%27, permblocks_dest)::Tuple
│         (biperm_dest = Core._apply_iterate(Base.iterate, %26, %28))
│         (permblocks1 = Core.tuple(perm_codomain1, perm_domain1))
│   %31 = NDTensors.TensorAlgebra.blockedperm::Core.Const(NDTensors.TensorAlgebra.blockedperm)
│   %32 = !NDTensors.TensorAlgebra.isempty::Core.Const(!isempty)
│   %33 = NDTensors.TensorAlgebra.filter(%32, permblocks1)::Tuple
│         (biperm1 = Core._apply_iterate(Base.iterate, %31, %33))
│         (permblocks2 = Core.tuple(perm_codomain2, perm_domain2))
│   %36 = NDTensors.TensorAlgebra.blockedperm::Core.Const(NDTensors.TensorAlgebra.blockedperm)
│   %37 = !NDTensors.TensorAlgebra.isempty::Core.Const(!isempty)
│   %38 = NDTensors.TensorAlgebra.filter(%37, permblocks2)::Tuple
│         (biperm2 = Core._apply_iterate(Base.iterate, %36, %38))
│   %40 = Core.tuple(biperm_dest, biperm1, biperm2)::Tuple{Any, Any, Any}
└──       return %40

It is called in contract at ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/contract.jl. If I directly feed the output of blockedperms to contract, the result is type stable:

biperm_dest = TensorAlgebra.blockedperm((1, 2), (3,))
biperm1 = TensorAlgebra.blockedperm((1, 2), (3,))
biperm2 = TensorAlgebra.blockedperm((1, ), (2,))
@code_warntype TensorAlgebra.contract(TensorAlgebra.default_contract_alg(), biperm_dest, ones((1,1,1)), biperm1, ones((1,1)), biperm2, true)
MethodInstance for NDTensors.TensorAlgebra.contract(::Algorithm{:matricize, @NamedTuple{}}, ::NDTensors.TensorAlgebra.BlockedPermutation{2, 3, Tuple{Tuple{Int64, Int64}, Tuple{Int64}}}, ::Array{Float64, 3}, ::NDTensors.TensorAlgebra.BlockedPermutation{2, 3, Tuple{Tuple{Int64, Int64}, Tuple{Int64}}}, ::Matrix{Float64}, ::NDTensors.TensorAlgebra.BlockedPermutation{2, 2, Tuple{Tuple{Int64}, Tuple{Int64}}}, ::Bool)
  from contract(alg::Algorithm, biperm_dest::NDTensors.TensorAlgebra.BlockedPermutation, a1::AbstractArray, biperm1::NDTensors.TensorAlgebra.BlockedPermutation, a2::AbstractArray, biperm2::NDTensors.TensorAlgebra.BlockedPermutation, α::Number; kwargs...) @ NDTensors.TensorAlgebra ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/contract.jl:107
Arguments
  #self#::Core.Const(NDTensors.TensorAlgebra.contract)
  alg::Core.Const(Algorithm type matricize, NamedTuple())
  biperm_dest::NDTensors.TensorAlgebra.BlockedPermutation{2, 3, Tuple{Tuple{Int64, Int64}, Tuple{Int64}}}
  a1::Array{Float64, 3}
  biperm1::NDTensors.TensorAlgebra.BlockedPermutation{2, 3, Tuple{Tuple{Int64, Int64}, Tuple{Int64}}}
  a2::Matrix{Float64}
  biperm2::NDTensors.TensorAlgebra.BlockedPermutation{2, 2, Tuple{Tuple{Int64}, Tuple{Int64}}}
  α::Bool
Body::Array{Float64, 3}
1 ─ %1 = Core.NamedTuple()::Core.Const(NamedTuple())
│   %2 = Base.pairs(%1)::Core.Const(Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}())
│   %3 = NDTensors.TensorAlgebra.:(var"#contract#37")(%2, #self#, alg, biperm_dest, a1, biperm1, a2, biperm2, α)::Array{Float64, 3}
└──      return %3

[ogauthe]

Pasted from ITensor/ITensors.jl#1507

For a FusionTensor, a permutation is not enough to specify the structure, a biperm is needed. Currently, TensorAlgebra.contract does not allow this.

using NDTensors.FusionTensors: FusionTensor
using NDTensors.GradedAxes
using NDTensors.Sectors: U1

  g1 = GradedAxes.gradedrange([U1(0) => 1])
  ft1 = FusionTensor(Float64, (GradedAxes.dual(g1),), (g1,))
  ft2 = FusionTensor(Float64, (GradedAxes.dual(g1),), (g1,g1))

  ft3, legs = TensorAlgebra.contract(ft1, (1,2), ft2, (2, 3, 4))  # Ok
  ft4 = TensorAlgebra.contract((1,4, 3), ft1, (1,2), ft2, (2, 3, 4))  # Ok
  ft5 = TensorAlgebra.contract(((1,4),(3,)), ft1, (1,2), ft2, (2, 3, 4))   # MethodError

ERROR: MethodError: no method matching blockedperm(::Tuple{Nothing}, ::Tuple{Nothing, Nothing})

Closest candidates are:
blockedperm(::NDTensors.NamedDimsArrays.AbstractNamedDimsArray, ::Tuple...)
 @ NDTensors ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/NamedDimsArrays/ext/NamedDimsArraysTensorAlgebraExt/src/fusedims.jl:4

Stacktrace:
[1] blockedperms(::typeof(NDTensors.TensorAlgebra.contract), dimnames_dest::Tuple{…}, dimnames1::Tuple{…}, dimnames2::Tuple{…})
 @ NDTensors.TensorAlgebra ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/blockedperms.jl:26
[2] contract(alg::NDTensors.BackendSelection.Algorithm{…}, labels_dest::Tuple{…}, a1::FusionTensor{…}, labels1::Tuple{…}, a2::FusionTensor{…}, labels2::Tuple{…}, α::Bool; kwargs::@Kwargs{})
 @ NDTensors.TensorAlgebra ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/contract.jl:87
[3] contract
 @ ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/contract.jl:77 [inlined]
[4] #contract#33
 @ ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/contract.jl:58 [inlined]
[5] contract
 @ ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/contract.jl:48 [inlined]
[6] contract(labels_dest::Tuple{…}, a1::FusionTensor{…}, labels1::Tuple{…}, a2::FusionTensor{…}, labels2::Tuple{…})
 @ NDTensors.TensorAlgebra ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/contract.jl:48
[7] top-level scope
 @ REPL[65]:1
Some type information was truncated. Use `show(err)` to see complete types.

Note that it is not possible to construct a BlockedPermutation from ((1,4), (3,)):

julia> TensorAlgebra.blockedperm((1,4), (3,))
ERROR: AssertionError: isperm(blockedperm)
Stacktrace:
[1] blockedperm
 @ ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/blockedpermutation.jl:166 [inlined]
[2] blockedperm
 @ ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/blockedpermutation.jl:110 [inlined]
[3] #blockedperm#13
 @ ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/blockedpermutation.jl:106 [inlined]
[4] blockedperm(::Tuple{Int64, Int64}, ::Tuple{Int64})
 @ NDTensors.TensorAlgebra ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/blockedpermutation.jl:105
[5] top-level scope
 @ REPL[66]:1

but even a correct BlockedPermutation fails:

biperm = TensorAlgebra.blockedperm((1,3),(2,))
ft5 = TensorAlgebra.contract(biperm, ft1, (1,4), ft2, (4, 2, 3))   # MethodError

ERROR: MethodError: no method matching contract(::NDTensors.TensorAlgebra.BlockedPermutation{…}, ::FusionTensor{…}, ::Tuple{…}, ::FusionTensor{…}, ::Tuple{…})

Closest candidates are:
  contract(::NDTensors.BackendSelection.Algorithm, ::AbstractArray, ::Tuple, ::AbstractArray, ::Tuple; ...)
   @ NDTensors ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/contract.jl:35
  contract(::NDTensors.BackendSelection.Algorithm, ::AbstractArray, ::Tuple, ::AbstractArray, ::Tuple, ::Number; kwargs...)
   @ NDTensors ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/contract.jl:35
  contract(::Tuple, ::AbstractArray, ::Tuple, ::AbstractArray, ::Tuple; ...)
   @ NDTensors ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/contract.jl:48
  ...

Stacktrace:
 [1] top-level scope
   @ REPL[84]:1
Some type information was truncated. Use `show(err)` to see complete types.

[ogauthe]

issue: contracting zero axis crashes. I think this should return the outer product of two arrays.

TensorAlgebra.contract( ones((2,)), (1,), ones((2,)), (2,))

ERROR: MethodError: no method matching output_axes(::typeof(contract), ::NDTensors.TensorAlgebra.BlockedPermutation{…}, ::Vector{…}, ::NDTensors.TensorAlgebra.BlockedPermutation{…}, ::Vector{…}, ::NDTensors.TensorAlgebra.BlockedPermutation{…}, ::Bool)

Closest candidates are:
  output_axes(::typeof(contract), ::NDTensors.TensorAlgebra.BlockedPermutation{0, Length, Blocks} where {Length, Blocks<:Tuple{}}, ::AbstractArray, ::NDTensors.TensorAlgebra.BlockedPermutation{1, Length, Blocks} where {Length, Blocks<:Tuple{Tuple{Vararg{Int64}}}}, ::AbstractArray, ::NDTensors.TensorAlgebra.BlockedPermutation{1, Length, Blocks} where {Length, Blocks<:Tuple{Tuple{Vararg{Int64}}}}, ::Number)
   @ NDTensors ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/allocate_output.jl:23
  output_axes(::typeof(contract), ::NDTensors.TensorAlgebra.BlockedPermutation{1, Length, Blocks} where {Length, Blocks<:Tuple{Tuple{Vararg{Int64}}}}, ::AbstractArray, ::NDTensors.TensorAlgebra.BlockedPermutation{2, Length, Blocks} where {Length, Blocks<:Tuple{Tuple{Vararg{Int64}}, Tuple{Vararg{Int64}}}}, ::AbstractArray, ::NDTensors.TensorAlgebra.BlockedPermutation{1, Length, Blocks} where {Length, Blocks<:Tuple{Tuple{Vararg{Int64}}}}, ::Number)
   @ NDTensors ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/allocate_output.jl:55
  output_axes(::typeof(contract), ::NDTensors.TensorAlgebra.BlockedPermutation{2, Length, Blocks} where {Length, Blocks<:Tuple{Tuple{Vararg{Int64}}, Tuple{Vararg{Int64}}}}, ::AbstractArray, ::NDTensors.TensorAlgebra.BlockedPermutation{2, Length, Blocks} where {Length, Blocks<:Tuple{Tuple{Vararg{Int64}}, Tuple{Vararg{Int64}}}}, ::AbstractArray, ::NDTensors.TensorAlgebra.BlockedPermutation{2, Length, Blocks} where {Length, Blocks<:Tuple{Tuple{Vararg{Int64}}, Tuple{Vararg{Int64}}}}, ::Number)
   @ NDTensors ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/allocate_output.jl:5
  ...

Stacktrace:
  [1] allocate_output(::typeof(contract), biperm_dest::NDTensors.TensorAlgebra.BlockedPermutation{…}, a1::Vector{…}, biperm1::NDTensors.TensorAlgebra.BlockedPermutation{…}, a2::Vector{…}, biperm2::NDTensors.TensorAlgebra.BlockedPermutation{…}, α::Bool)
    @ NDTensors.TensorAlgebra ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/allocate_output.jl:81
  [2] contract(alg::NDTensors.BackendSelection.Algorithm{…}, biperm_dest::NDTensors.TensorAlgebra.BlockedPermutation{…}, a1::Vector{…}, biperm1::NDTensors.TensorAlgebra.BlockedPermutation{…}, a2::Vector{…}, biperm2::NDTensors.TensorAlgebra.BlockedPermutation{…}, α::Bool; kwargs::@Kwargs{})
    @ NDTensors.TensorAlgebra ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/contract.jl:117
  [3] contract(alg::NDTensors.BackendSelection.Algorithm{…}, biperm_dest::NDTensors.TensorAlgebra.BlockedPermutation{…}, a1::Vector{…}, biperm1::NDTensors.TensorAlgebra.BlockedPermutation{…}, a2::Vector{…}, biperm2::NDTensors.TensorAlgebra.BlockedPermutation{…}, α::Bool)
    @ NDTensors.TensorAlgebra ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/contract.jl:107
  [4] contract(alg::NDTensors.BackendSelection.Algorithm{…}, labels_dest::Tuple{…}, a1::Vector{…}, labels1::Tuple{…}, a2::Vector{…}, labels2::Tuple{…}, α::Bool; kwargs::@Kwargs{})
    @ NDTensors.TensorAlgebra ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/contract.jl:88
  [5] contract(alg::NDTensors.BackendSelection.Algorithm{…}, labels_dest::Tuple{…}, a1::Vector{…}, labels1::Tuple{…}, a2::Vector{…}, labels2::Tuple{…}, α::Bool)
    @ NDTensors.TensorAlgebra ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/contract.jl:77
  [6] contract(alg::NDTensors.BackendSelection.Algorithm{…}, a1::Vector{…}, labels1::Tuple{…}, a2::Vector{…}, labels2::Tuple{…}, α::Bool; kwargs::@Kwargs{})
    @ NDTensors.TensorAlgebra ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/contract.jl:45
  [7] contract
    @ ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/contract.jl:35 [inlined]
  [8] #contract#31
    @ ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/contract.jl:32 [inlined]
  [9] contract
    @ ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/contract.jl:23 [inlined]
 [10] contract(a1::Vector{Float64}, labels1::Tuple{Int64}, a2::Vector{Float64}, labels2::Tuple{Int64})
    @ NDTensors.TensorAlgebra ~/Documents/itensor/ITensors.jl/NDTensors/src/lib/TensorAlgebra/src/contract/contract.jl:23
 [11] top-level scope
    @ REPL[28]:1
Some type information was truncated. Use `show(err)` to see complete types.

[mtfishman]

Pasted from ITensor/ITensors.jl#1507

For a FusionTensor, a permutation is not enough to specify the structure, a biperm is needed. Currently, TensorAlgebra.contract does not allow this.

I think this is supported, but you have to explicitly use the BlockedPermutation{2} type, i.e.: https://github.com/ITensor/ITensors.jl/blob/v0.6.23/NDTensors/src/lib/TensorAlgebra/src/contract/contract.jl#L107-L120.

EDIT: I see you tried that and faced some issues. Insetead of:

biperm = TensorAlgebra.blockedperm((1,3),(2,))
ft5 = TensorAlgebra.contract(biperm, ft1, (1,4), ft2, (4, 2, 3))   # MethodErro

you have to input all of them as BlockedPermutation, you can't mix using labels and blocked permutations. That seems like it would make the code pretty complicated. Also, it isn't clear to me if your biperm is really meant to be a permutation or a set of labels.

I think BlockedPermutation isn't really what you're looking for, and we maybe need another type, like a BlockedTuple, which indicates the tuple is split into parts (so BlockedTuple{2} would indicated it is a set of labels split into two sets representing the codomain and domain).

[mtfishman]

@ogauthe here is an idea for supporting labels that are split into codomain and domain. We can define a BlockedTuple{N} type, which would be like the BlockedPermutation{N} (and other blocked types in BlockArrays) that allows splitting up a Tuple into N blocks. (In fact, the BlockedPermutation{N} type could just be a special kind of BlockedTuple{N} that happens to be a permutation.) Then, the code logic for contract could have another layer of functions that can accept BlockedTuple{2} as blocked labels that specify the partitioning, which can then be compiled down to BlockedPermutation{2} and the that layer of code can remain the same.

[ogauthe]

Indeed BlockedPermutation does not seems suited for this. We could use TensorAlgebra.TupleOfTuples, defined as const TupleOfTuples{N} = Tuple{Vararg{Tuple{Vararg{Int}},N}}

EDIT: as it is just an alias for a Base type, we should avoid defining a BlockArrays interface for it, but I think it still fits the purpose.

[mtfishman]

That's definitely a possibility, as opposed to defining a devoted type BlockedTuple. It's a design I've gone back and forth in my head about, because I want to be as general as possible about what the labels might be, i.e. what you wrote only allows Int labels, while we want to allow more general labels (for example, we want to be able to use the ITensors.Index type as a label).

So the issue with just using a Tuple of Tuples instead of a devoted BlockedTuple type is, how do we distinguish that it isn't a flat set of labels which happen to be Tuples, rather than indicating a partitioning of the labels? We could maybe do that based on the number of labels, but I think that becomes too subtle and will make the code hard to design (for example, there is an ambiguity in the case of matrix inputs).

[mtfishman]

Either way, I think I would want to have a BlockedTuple type that is used inside the code, and maybe have a compatibility layer on top of that where, if a Tuple of Tuples is input as the labels, we assume that is a nested structure and convert it to a BlockedTuple. I think that will make coding all of this much more sane.