fix tensor_product for zero and one blocklength axes (#13)

Author: ogauthe

Project.toml

@@ -1,7 +1,7 @@
 name = "TensorProducts"
 uuid = "decf83d6-1968-43f4-96dc-fdb3fe15fc6d"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.1.4"
+version = "0.1.5"
 
 [weakdeps]
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"

ext/TensorProductsBlockArraysExt/TensorProductsBlockArraysExt.jl

@@ -14,9 +14,9 @@ using TensorProducts: OneToOne, TensorProducts
 function TensorProducts.tensor_product(
   a1::AbstractBlockedUnitRange, a2::AbstractBlockedUnitRange
 )
-  new_blocklengths = mapreduce(vcat, Iterators.product(blocks(a1), blocks(a2))) do (x, y)
-    return length(x) * length(y)
-  end
+  new_blocklengths = vec(
+    map(splat(*), Iterators.product(blocklengths(a1), blocklengths(a2)))
+  )
   return blockedrange(new_blocklengths)
 end
 

test/test_tensor_product.jl

@@ -5,7 +5,9 @@ using TensorProducts: ⊗, OneToOne, tensor_product
 using BlockArrays: blockedrange, blockisequal
 
 r0 = OneToOne()
-b1 = blockedrange([1, 2])
+b0 = blockedrange(Int[])
+b1 = blockedrange([1])
+b2 = blockedrange([1, 2])
 
 @testset "tensor_product" begin
   @test tensor_product() isa OneToOne
@@ -21,8 +23,13 @@ b1 = blockedrange([1, 2])
   @test blockisequal(tensor_product(b1, r0), b1)
   @test blockisequal(tensor_product(r0, b1), b1)
 
-  @test blockisequal(tensor_product(b1, b1), blockedrange([1, 2, 2, 4]))
-  @test blockisequal(tensor_product(b1, b1), blockedrange([1, 2, 2, 4]))
+  @test blockisequal(tensor_product(b0, b0), b0)
+  @test blockisequal(tensor_product(b0, b1), b0)
+  @test blockisequal(tensor_product(b1, b0), b0)
+  @test blockisequal(tensor_product(b1, b1), b1)
+  @test blockisequal(tensor_product(b1, b2), b2)
+  @test blockisequal(tensor_product(b2, b1), b2)
+  @test blockisequal(tensor_product(b2, b2), blockedrange([1, 2, 2, 4]))
 
   @test ⊗(r0, r0) isa OneToOne
 end