KroneckerArraysTensorProductsExt

Author: mtfishman

Project.toml

@@ -16,9 +16,11 @@ MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
 [weakdeps]
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 BlockSparseArrays = "2c9a651f-6452-4ace-a6ac-809f4280fbb4"
+TensorProducts = "decf83d6-1968-43f4-96dc-fdb3fe15fc6d"
 
 [extensions]
 KroneckerArraysBlockSparseArraysExt = ["BlockArrays", "BlockSparseArrays"]
+KroneckerArraysTensorProductsExt = "TensorProducts"
 
 [compat]
 Adapt = "4.3"
@@ -31,4 +33,5 @@ GPUArraysCore = "0.2"
 LinearAlgebra = "1.10"
 MapBroadcast = "0.1.9"
 MatrixAlgebraKit = "0.2"
+TensorProducts = "0.1.7"
 julia = "1.10"

ext/KroneckerArraysTensorProductsExt/KroneckerArraysTensorProductsExt.jl

@@ -1,11 +1,11 @@
 module KroneckerArraysTensorProductsExt
 
-using KroneckerArrays: CartesianProductOneTo, ×, arg1, arg2, cartesianrange
+using KroneckerArrays: CartesianProductOneTo, ×, arg1, arg2, cartesianrange, unproduct
 using TensorProducts: TensorProducts, tensor_product
 function TensorProducts.tensor_product(a1::CartesianProductOneTo, a2::CartesianProductOneTo)
-  return cartesianrange(
-    tensor_product(arg1(a1), arg1(a2)) × tensor_product(arg2(a1), arg2(a2))
-  )
+  prod = tensor_product(arg1(a1), arg1(a2)) × tensor_product(arg2(a1), arg2(a2))
+  range = tensor_product(unproduct(a1), unproduct(a2))
+  return cartesianrange(prod, range)
 end
 
 end

test/Project.toml

@@ -14,6 +14,7 @@ MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
+TensorProducts = "decf83d6-1968-43f4-96dc-fdb3fe15fc6d"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a"
 
@@ -33,5 +34,6 @@ MatrixAlgebraKit = "0.2"
 SafeTestsets = "0.1"
 StableRNGs = "1.0"
 Suppressor = "0.2"
+TensorProducts = "0.1.7"
 Test = "1.10"
 TestExtras = "0.3"

test/test_tensorproducts.jl

@@ -0,0 +1,13 @@
+using KroneckerArrays: ×, arg1, arg2, cartesianrange, unproduct
+using TensorProducts: tensor_product
+using Test: @test, @testset
+
+@testset "KroneckerArraysTensorProductsExt" begin
+  r1 = cartesianrange(2, 3)
+  r2 = cartesianrange(4, 5)
+  r = tensor_product(r1, r2)
+  @test r ≡ cartesianrange(8, 15)
+  @test arg1(r) ≡ Base.OneTo(8)
+  @test arg2(r) ≡ Base.OneTo(15)
+  @test unproduct(r) ≡ Base.OneTo(120)
+end