rewrite with invbiperm. BlockArray fails.

Author: ogauthe

src/contract/allocate_output.jl

@@ -17,7 +17,7 @@ end
 # i.e. `ContractAdd`?
 function output_axes(
   ::typeof(contract),
-  biperm_out::AbstractBlockPermutation{2},
+  biperm_a12_to_dest::AbstractBlockPermutation{2},
   a1::AbstractArray,
   biperm1::AbstractBlockPermutation{2},
   a2::AbstractArray,
@@ -28,21 +28,21 @@ function output_axes(
   axes_contracted2, axes_domain = blocks(axes(a2)[biperm2])
   @assert axes_contracted == axes_contracted2
   # default: flatten biperm_out
-  return genperm((axes_codomain..., axes_domain...), Tuple(biperm_out))
+  return genperm((axes_codomain..., axes_domain...), Tuple(biperm_a12_to_dest))
 end
 
 # TODO: Use `ArrayLayouts`-like `MulAdd` object,
 # i.e. `ContractAdd`?
 function allocate_output(
   ::typeof(contract),
-  biperm_out::AbstractBlockPermutation,
+  biperm_a12_to_dest::AbstractBlockPermutation,
   a1::AbstractArray,
   biperm1::AbstractBlockPermutation,
   a2::AbstractArray,
   biperm2::AbstractBlockPermutation,
   α::Number=one(Bool),
 )
   check_input(contract, a1, biperm1, a2, biperm2)
-  axes_dest = output_axes(contract, biperm_out, a1, biperm1, a2, biperm2, α)
+  axes_dest = output_axes(contract, biperm_a12_to_dest, a1, biperm1, a2, biperm2, α)
   return similar(a1, promote_type(eltype(a1), eltype(a2), typeof(α)), axes_dest)
 end

src/contract/blockedperms.jl

@@ -1,25 +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
+# default: if no bipartion is specified, all axes to domain
+invbiperm(perm, ::Any) = invbiperm(perm, Val(0))
+invbiperm(perm, t::Tuple{Tuple,Tuple}) = invbiperm(perm, tuplemortar(t))
+invbiperm(perm, t::AbstractBlockTuple{2}) = invbiperm(perm, Val(first(blocklength(t))))
 
 function invbiperm(perm, ::Val{N1}) where {N1}
   perm_out = invperm(Tuple(perm))
+  length(perm) <= N1 && return blockedpermvcat(perm_out, ())
   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)
+function blockedperms(
+  f::typeof(contract), alg::Algorithm, dimnames_dest, dimnames1, dimnames2
+)
+  return blockedperms(f, dimnames_dest, dimnames1, dimnames2)
 end
 
-function blockedperms(::typeof(contract), dimnames_dest::BlockedTuple, dimnames1, dimnames2)
+# codomain <-- domain
+function blockedperms(::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"))
@@ -31,9 +31,8 @@ function blockedperms(::typeof(contract), dimnames_dest::BlockedTuple, dimnames1
 
   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)))
-  )
+  biperm_dest_to_a12 = (perm_codomain_dest..., perm_domain_dest...)
+  biperm_a12_to_dest = invbiperm(biperm_dest_to_a12, dimnames_dest)
 
   perm_codomain1 = BaseExtensions.indexin(codomain, dimnames1)
   perm_domain1 = BaseExtensions.indexin(contracted, dimnames1)
@@ -45,5 +44,5 @@ function blockedperms(::typeof(contract), dimnames_dest::BlockedTuple, dimnames1
   biperm1 = blockedpermvcat(permblocks1...)
   permblocks2 = (perm_codomain2, perm_domain2)
   biperm2 = blockedpermvcat(permblocks2...)
-  return biperm_out, biperm1, biperm2
+  return biperm_a12_to_dest, biperm1, biperm2
 end

src/contract/contract.jl

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

src/contract/contract_matricize/contract.jl

@@ -1,63 +1,21 @@
 using LinearAlgebra: mul!
 
-function isinplace(::AbstractArray, biperm_out)
-  return istrivialperm(Tuple(biperm_out))
-end
-
 function contract!(
-  alg::Matricize,
-  a_dest::AbstractArray,
-  biperm_out::AbstractBlockPermutation{2},
-  a1::AbstractArray,
-  biperm1::AbstractBlockPermutation{2},
-  a2::AbstractArray,
-  biperm2::AbstractBlockPermutation{2},
-  α::Number,
-  β::Number,
-)
-  if isinplace(a_dest, biperm_out)
-    return contract_inplace!(alg, a_dest, biperm_out, a1, biperm1, a2, biperm2, α, β)
-  else
-    return contract_outofplace!(alg, a_dest, biperm_out, a1, biperm1, a2, biperm2, α, β)
-  end
-end
-
-function contract_inplace!(
-  ::Matricize,
-  a_dest::AbstractArray,
-  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; copy=false)
-  a1_mat = matricize(a1, biperm1; copy=false)
-  a2_mat = matricize(a2, biperm2; copy=false)
-  mul!(a_dest_mat, a1_mat, a2_mat, α, β)
-  return a_dest
-end
-
-function contract_outofplace!(
   ::Matricize,
   a_dest::AbstractArray,
-  biperm_out::AbstractBlockPermutation{2},
+  biperm_a12_to_dest::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)
+  biperm_dest_to_a12 = invbiperm(biperm_a12_to_dest, Val(first(blocklengths(biperm1))))
+  check_input(contract, a_dest, biperm_dest_to_a12, a1, biperm1, a2, biperm2)
   a1_mat = matricize(a1, biperm1)
   a2_mat = matricize(a2, biperm2)
   a_dest_mat = a1_mat * a2_mat
-  unmatricize_add!(a_dest, a_dest_mat, biperm_dest, α, β)
+  unmatricize_add!(a_dest, a_dest_mat, biperm_dest_to_a12, α, β)
   return a_dest
 end

src/matricize.jl

@@ -45,23 +45,20 @@ end
 # matrix factorizations assume copy
 # maybe: copy=false kwarg
 
-function matricize(a::AbstractArray, biperm::AbstractBlockPermutation{2}; copy=false)
+function matricize(a::AbstractArray, biperm::AbstractBlockPermutation{2})
   ndims(a) == length(biperm) || throw(ArgumentError("Invalid bipermutation"))
-  return matricize(FusionStyle(a), a, biperm; copy)
+  return matricize(FusionStyle(a), a, biperm)
 end
 
 function matricize(
-  style::FusionStyle, a::AbstractArray, biperm::AbstractBlockPermutation{2}; copy=false
+  style::FusionStyle, a::AbstractArray, biperm::AbstractBlockPermutation{2}
 )
-  if istrivialperm(Tuple(biperm)) && !copy
-    return matricize(style, a, trivialperm(biperm))
-  end
   a_perm = permuteblockeddims(a, biperm)
   return matricize(style, a_perm, trivialperm(biperm))
 end
 
 function matricize(
-  style::FusionStyle, a::AbstractArray, biperm::BlockedTrivialPermutation{2}; copy=false
+  style::FusionStyle, a::AbstractArray, biperm::BlockedTrivialPermutation{2}
 )
   return throw(MethodError(matricize, Tuple{typeof(style),typeof(a),typeof(biperm)}))
 end
@@ -72,22 +69,29 @@ function matricize(::ReshapeFusion, a::AbstractArray, biperm::BlockedTrivialPerm
   return reshape(a, new_axes...)
 end
 
-function matricize(a::AbstractArray, permblock1::Tuple, permblock2::Tuple; copy=false)
-  return matricize(a, blockedpermvcat(permblock1, permblock2; length=Val(ndims(a))); copy)
+function matricize(a::AbstractArray, permblock1::Tuple, permblock2::Tuple)
+  return matricize(a, blockedpermvcat(permblock1, permblock2; length=Val(ndims(a))))
 end
 
 # ====================================  unmatricize  =======================================
-function unmatricize(m::AbstractMatrix, axes, biperm::AbstractBlockPermutation{2})
-  length(axes) == length(biperm) || throw(ArgumentError("axes do not match permutation"))
-  return unmatricize(FusionStyle(m), m, axes, biperm)
+function unmatricize(
+  m::AbstractMatrix, axes_dest, biperm_dest_to_a12::AbstractBlockPermutation{2}
+)
+  length(axes_dest) == length(biperm_dest_to_a12) ||
+    throw(ArgumentError("axes do not match permutation"))
+  return unmatricize(FusionStyle(m), m, axes_dest, biperm_dest_to_a12)
 end
 
 function unmatricize(
-  ::FusionStyle, m::AbstractMatrix, axes, biperm::AbstractBlockPermutation{2}
+  ::FusionStyle,
+  m::AbstractMatrix,
+  axes_dest,
+  biperm_dest_to_a12::AbstractBlockPermutation{2},
 )
-  blocked_axes = axes[biperm]
-  a_perm = unmatricize(m, blocked_axes)
-  return permuteblockeddims(a_perm, invperm(biperm))
+  blocked_axes = axes_dest[biperm_dest_to_a12]
+  a12 = unmatricize(m, blocked_axes)
+  biperm_a12_to_dest = invbiperm(biperm_dest_to_a12, axes_dest)
+  return permuteblockeddims(a12, biperm_a12_to_dest)
 end
 
 function unmatricize(
@@ -111,14 +115,19 @@ function unmatricize(
   return unmatricize(m, blocked_axes)
 end
 
-function unmatricize!(a, m::AbstractMatrix, biperm::AbstractBlockPermutation{2})
-  ndims(a) == length(biperm) ||
+function unmatricize!(
+  a_dest, m::AbstractMatrix, biperm_dest_to_a12::AbstractBlockPermutation{2}
+)
+  ndims(a_dest) == length(biperm_dest_to_a12) ||
     throw(ArgumentError("destination does not match permutation"))
-  blocked_axes = axes(a)[biperm]
+  blocked_axes = axes(a_dest)[biperm_dest_to_a12]
   a_perm = unmatricize(m, blocked_axes)
-  return permuteblockeddims!(a, a_perm, invperm(biperm))
+  biperm_a12_to_dest = invbiperm(biperm_dest_to_a12, axes(a_dest))
+  return permuteblockeddims!(a_dest, a_perm, biperm_a12_to_dest)
 end
 
-function unmatricize_add!(a_dest, a_dest_mat, biperm_dest, α, β)
-  return mul!(a_dest, 1.0, unmatricize(a_dest_mat, axes(a_dest), biperm_dest), α, β)
+function unmatricize_add!(a_dest, a_dest_mat, biperm_dest_to_a12, α, β)
+  a12 = unmatricize(a_dest_mat, axes(a_dest), biperm_dest_to_a12)
+  a_dest .= α .* a12 .+ β .* a_dest
+  return a_dest
 end

test/test_basics.jl

@@ -95,9 +95,10 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
     @test a ≈ a0
 
     bp = blockedpermvcat((4, 2), (1, 3))
-    a = unmatricize(m, map(i -> axes0[i], invperm(Tuple(bp))), bp)
+    bpinv = blockedpermvcat((3, 2), (4, 1))
+    a = unmatricize(m, map(i -> axes0[i], bp), bpinv)
     @test eltype(a) === elt
-    @test a ≈ permutedims(a0, invperm(Tuple(bp)))
+    @test a ≈ permutedims(a0, Tuple(bp))
 
     a = similar(a0)
     unmatricize!(a, m, blockedpermvcat((1, 2), (3, 4)))
@@ -109,7 +110,7 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
 
     a1 = permutedims(a0, Tuple(bp))
     a = similar(a1)
-    unmatricize!(a, m, invperm(bp))
+    unmatricize!(a, m, bpinv)
     @test a ≈ a1
 
     a = unmatricize(m, (), axes0)