pass tests

Author: ogauthe

src/contract/allocate_output.jl

@@ -17,7 +17,7 @@ end
 # i.e. `ContractAdd`?
 function output_axes(
   ::typeof(contract),
-  biperm_dest::AbstractBlockPermutation{2},
+  biperm_out::AbstractBlockPermutation{2},
   a1::AbstractArray,
   biperm1::AbstractBlockPermutation{2},
   a2::AbstractArray,
@@ -27,23 +27,22 @@ function output_axes(
   axes_codomain, axes_contracted = blocks(axes(a1)[biperm1])
   axes_contracted2, axes_domain = blocks(axes(a2)[biperm2])
   @assert axes_contracted == axes_contracted2
-  return genperm((axes_codomain..., axes_domain...), invperm(Tuple(biperm_dest)))
+  # default: flatten biperm_out
+  return genperm((axes_codomain..., axes_domain...), Tuple(biperm_out))
 end
 
 # TODO: Use `ArrayLayouts`-like `MulAdd` object,
 # i.e. `ContractAdd`?
 function allocate_output(
   ::typeof(contract),
-  biperm_dest::AbstractBlockPermutation,
+  biperm_out::AbstractBlockPermutation,
   a1::AbstractArray,
   biperm1::AbstractBlockPermutation,
   a2::AbstractArray,
   biperm2::AbstractBlockPermutation,
   α::Number=one(Bool),
 )
   check_input(contract, a1, biperm1, a2, biperm2)
-  blocklengths(biperm_dest) == (length(biperm1[Block(1)]), length(biperm2[Block(2)])) ||
-    throw(ArgumentError("Invalid permutation for destination tensor"))
-  axes_dest = output_axes(contract, biperm_dest, a1, biperm1, a2, biperm2, α)
+  axes_dest = output_axes(contract, biperm_out, a1, biperm1, a2, biperm2, α)
   return similar(a1, promote_type(eltype(a1), eltype(a2), typeof(α)), axes_dest)
 end

src/contract/blockedperms.jl

@@ -1,13 +1,25 @@
 using .BaseExtensions: BaseExtensions
+using BlockArrays: blocklengths
 
 function blockedperms(
   f::typeof(contract), alg::Algorithm, dimnames_dest, dimnames1, dimnames2
 )
   return blockedperms(f, dimnames_dest, dimnames1, dimnames2)
 end
 
+function invbiperm(perm, ::Val{N1}) where {N1}
+  perm_out = invperm(Tuple(perm))
+  return blockedpermvcat(perm_out[begin:N1], (perm_out[(N1 + 1):end]))
+end
+
 # codomain <-- domain
 function blockedperms(::typeof(contract), dimnames_dest, dimnames1, dimnames2)
+  # default: if no bipartion is specified, all axes to domain
+  dimnames_dest_bt = tuplemortar(((), Tuple(dimnames_dest)))
+  return blockedperms(contract, dimnames_dest_bt, dimnames1, dimnames2)
+end
+
+function blockedperms(::typeof(contract), dimnames_dest::BlockedTuple, dimnames1, dimnames2)
   dimnames = collect(Iterators.flatten((dimnames_dest, dimnames1, dimnames2)))
   for i in unique(dimnames)
     count(==(i), dimnames) == 2 || throw(ArgumentError("Invalid contraction labels"))
@@ -19,18 +31,19 @@ function blockedperms(::typeof(contract), dimnames_dest, dimnames1, dimnames2)
 
   perm_codomain_dest = BaseExtensions.indexin(codomain, dimnames_dest)
   perm_domain_dest = BaseExtensions.indexin(domain, dimnames_dest)
+  biperm_out = invbiperm(
+    (perm_codomain_dest..., perm_domain_dest...), Val(first(blocklengths(dimnames_dest)))
+  )
 
   perm_codomain1 = BaseExtensions.indexin(codomain, dimnames1)
   perm_domain1 = BaseExtensions.indexin(contracted, dimnames1)
 
   perm_codomain2 = BaseExtensions.indexin(contracted, dimnames2)
   perm_domain2 = BaseExtensions.indexin(domain, dimnames2)
 
-  permblocks_dest = (perm_codomain_dest, perm_domain_dest)
-  biperm_dest = blockedpermvcat(permblocks_dest...)
   permblocks1 = (perm_codomain1, perm_domain1)
   biperm1 = blockedpermvcat(permblocks1...)
   permblocks2 = (perm_codomain2, perm_domain2)
   biperm2 = blockedpermvcat(permblocks2...)
-  return biperm_dest, biperm1, biperm2
+  return biperm_out, biperm1, biperm2
 end

src/contract/contract.jl

@@ -13,7 +13,7 @@ default_contract_alg() = Matricize()
 function contract!(
   alg::Algorithm,
   a_dest::AbstractArray,
-  biperm_dest::AbstractBlockPermutation,
+  biperm_out::AbstractBlockPermutation,
   a1::AbstractArray,
   biperm1::AbstractBlockPermutation,
   a2::AbstractArray,
@@ -89,8 +89,8 @@ function contract(
   kwargs...,
 )
   check_input(contract, a1, labels1, a2, labels2)
-  biperm_dest, biperm1, biperm2 = blockedperms(contract, labels_dest, labels1, labels2)
-  return contract(alg, biperm_dest, a1, biperm1, a2, biperm2, α; kwargs...)
+  biperm_out, biperm1, biperm2 = blockedperms(contract, labels_dest, labels1, labels2)
+  return contract(alg, biperm_out, a1, biperm1, a2, biperm2, α; kwargs...)
 end
 
 function contract!(
@@ -106,13 +106,13 @@ function contract!(
   kwargs...,
 )
   check_input(contract, a_dest, labels_dest, a1, labels1, a2, labels2)
-  biperm_dest, biperm1, biperm2 = blockedperms(contract, labels_dest, labels1, labels2)
-  return contract!(alg, a_dest, biperm_dest, a1, biperm1, a2, biperm2, α, β; kwargs...)
+  biperm_out, biperm1, biperm2 = blockedperms(contract, labels_dest, labels1, labels2)
+  return contract!(alg, a_dest, biperm_out, a1, biperm1, a2, biperm2, α, β; kwargs...)
 end
 
 function contract(
   alg::Algorithm,
-  biperm_dest::AbstractBlockPermutation,
+  biperm_out::AbstractBlockPermutation,
   a1::AbstractArray,
   biperm1::AbstractBlockPermutation,
   a2::AbstractArray,
@@ -121,7 +121,7 @@ function contract(
   kwargs...,
 )
   check_input(contract, a1, biperm1, a2, biperm2)
-  a_dest = allocate_output(contract, biperm_dest, a1, biperm1, a2, biperm2, α)
-  contract!(alg, a_dest, biperm_dest, a1, biperm1, a2, biperm2, α, zero(Bool); kwargs...)
+  a_dest = allocate_output(contract, biperm_out, a1, biperm1, a2, biperm2, α)
+  contract!(alg, a_dest, biperm_out, a1, biperm1, a2, biperm2, α, zero(Bool); kwargs...)
   return a_dest
 end

src/contract/contract_matricize/contract.jl

@@ -3,14 +3,15 @@ using LinearAlgebra: mul!
 function contract!(
   ::Matricize,
   a_dest::AbstractArray,
-  biperm_dest::AbstractBlockPermutation{2},
+  biperm_out::AbstractBlockPermutation{2},
   a1::AbstractArray,
   biperm1::AbstractBlockPermutation{2},
   a2::AbstractArray,
   biperm2::AbstractBlockPermutation{2},
   α::Number,
   β::Number,
 )
+  biperm_dest = invbiperm(biperm_out, Val(first(blocklengths(biperm1))))
   check_input(contract, a_dest, biperm_dest, a1, biperm1, a2, biperm2)
   a_dest_mat = matricize(a_dest, biperm_dest)
   a1_mat = matricize(a1, biperm1)