define length_domain

Author: ogauthe

src/contract/allocate_output.jl

@@ -17,7 +17,7 @@ end
 # i.e. `ContractAdd`?
 function output_axes(
   ::typeof(contract),
-  biperm_a12_to_dest::AbstractBlockPermutation{2},
+  biperm_dest::AbstractBlockPermutation{2},
   a1::AbstractArray,
   biperm1::AbstractBlockPermutation{2},
   a2::AbstractArray,
@@ -26,23 +26,23 @@ function output_axes(
 )
   axes_codomain, axes_contracted = blocks(axes(a1)[biperm1])
   axes_contracted2, axes_domain = blocks(axes(a2)[biperm2])
-  @assert axes_contracted == axes_contracted2
+  @assert length.(axes_contracted) == length.(axes_contracted2)
   # default: flatten biperm_out
-  return genperm((axes_codomain..., axes_domain...), Tuple(biperm_a12_to_dest))
+  return genperm((axes_codomain..., axes_domain...), Tuple(biperm_dest))
 end
 
 # TODO: Use `ArrayLayouts`-like `MulAdd` object,
 # i.e. `ContractAdd`?
 function allocate_output(
   ::typeof(contract),
-  biperm_a12_to_dest::AbstractBlockPermutation,
+  biperm_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_a12_to_dest, a1, biperm1, a2, biperm2, α)
+  axes_dest = output_axes(contract, biperm_dest, a1, biperm1, a2, biperm2, α)
   return similar(a1, promote_type(eltype(a1), eltype(a2), typeof(α)), axes_dest)
 end

src/contract/blockedperms.jl

@@ -14,6 +14,8 @@ length_codomain(t::AbstractBlockTuple{2}) = first(blocklengths(t))
 # Assume all dimensions are in the domain by default
 length_codomain(t) = 0
 
+length_domain(t) = length(t) - length_codomain(t)
+
 function blockedperms(
   f::typeof(contract), alg::Algorithm, dimnames_dest, dimnames1, dimnames2
 )

src/matricize.jl

@@ -119,7 +119,6 @@ function unmatricize!(a_dest, m::AbstractMatrix, invbiperm::AbstractBlockPermuta
   blocked_axes = axes(a_dest)[invbiperm]
   a_perm = unmatricize(m, blocked_axes)
   biperm_dest = biperm(invperm(invbiperm), length_codomain(axes(a_dest)))
-
   return permuteblockeddims!(a_dest, a_perm, biperm_dest)
 end
 

test/test_basics.jl

@@ -12,6 +12,8 @@ using TensorAlgebra:
   contract,
   contract!,
   matricize,
+  length_codomain,
+  length_domain,
   tuplemortar,
   unmatricize,
   unmatricize!
@@ -20,6 +22,15 @@ default_rtol(elt::Type) = 10^(0.75 * log10(eps(real(elt))))
 const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
 
 @testset "TensorAlgebra" begin
+  @testset "misc" begin
+    t = (1, 2, 3)
+    bt = tuplemortar(((1, 2), (3,)))
+    @test length_codomain(t) == 0
+    @test length_codomain(bt) == 2
+    @test length_domain(t) == 3
+    @test length_domain(bt) == 1
+  end
+
   @testset "permuteblockeddims (eltype=$elt)" for elt in elts
     a = randn(elt, 2, 3, 4, 5)
     a_perm = permuteblockeddims(a, blockedpermvcat((3, 1), (2, 4)))