KroneckerArraysTensorAlgebraExt extension

Author: mtfishman

Project.toml

@@ -12,15 +12,16 @@ GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MapBroadcast = "ebd9b9da-f48d-417c-9660-449667d60261"
 MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
-TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 
 [weakdeps]
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 BlockSparseArrays = "2c9a651f-6452-4ace-a6ac-809f4280fbb4"
+TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 TensorProducts = "decf83d6-1968-43f4-96dc-fdb3fe15fc6d"
 
 [extensions]
 KroneckerArraysBlockSparseArraysExt = ["BlockArrays", "BlockSparseArrays"]
+KroneckerArraysTensorAlgebraExt = "TensorAlgebra"
 KroneckerArraysTensorProductsExt = "TensorProducts"
 
 [compat]

ext/KroneckerArraysTensorAlgebraExt/KroneckerArraysTensorAlgebraExt.jl

@@ -0,0 +1,41 @@
+module KroneckerArraysTensorAlgebraExt
+
+using KroneckerArrays: KroneckerArrays, KroneckerArray
+using TensorAlgebra: TensorAlgebra, AbstractBlockPermutation, FusionStyle, matricize
+
+struct KroneckerFusion{A<:FusionStyle,B<:FusionStyle} <: FusionStyle
+  a::A
+  b::B
+end
+KroneckerArrays.arg1(style::KroneckerFusion) = style.a
+KroneckerArrays.arg2(style::KroneckerFusion) = style.b
+function TensorAlgebra.FusionStyle(a::KroneckerArray)
+  return KroneckerFusion(FusionStyle(arg1(a)), FusionStyle(arg2(a)))
+end
+function matricize_kronecker(
+  style::KroneckerFusion, a::AbstractArray, biperm::AbstractBlockPermutation{2}
+)
+  return matricize(arg1(style), arg1(a), biperm) ⊗ matricize(arg2(style), arg2(a), biperm)
+end
+function TensorAlgebra.matricize(
+  style::KroneckerFusion, a::AbstractArray, biperm::AbstractBlockPermutation{2}
+)
+  return matricize_kronecker(style, a, biperm)
+end
+# Fix ambiguity error.
+# TODO: Investigate rewriting the logic in `TensorAlgebra.jl` to avoid this.
+using TensorAlgebra: BlockedTrivialPermutation, unmatricize
+function TensorAlgebra.matricize(
+  style::KroneckerFusion, a::AbstractArray, biperm::BlockedTrivialPermutation{2}
+)
+  return matricize_kronecker(style, a, biperm)
+end
+function unmatricize_kronecker(style::KroneckerFusion, a::AbstractArray, ax)
+  return unmatricize(arg1(style), arg1(a), arg1.(ax)) ⊗
+         unmatricize(arg2(style), arg2(a), arg2.(ax))
+end
+function TensorAlgebra.unmatricize(style::KroneckerFusion, a::AbstractArray, ax)
+  return unmatricize_kronecker(style, a, ax)
+end
+
+end

src/kroneckerarray.jl

@@ -494,39 +494,3 @@ function Base.broadcasted(
 )
   return broadcasted(style, /, a, f.args[2])
 end
-
-using TensorAlgebra: TensorAlgebra, AbstractBlockPermutation, FusionStyle, matricize
-struct KroneckerFusion{A<:FusionStyle,B<:FusionStyle} <: FusionStyle
-  a::A
-  b::B
-end
-arg1(style::KroneckerFusion) = style.a
-arg2(style::KroneckerFusion) = style.b
-function TensorAlgebra.FusionStyle(a::KroneckerArray)
-  return KroneckerFusion(FusionStyle(arg1(a)), FusionStyle(arg2(a)))
-end
-function matricize_kronecker(
-  style::KroneckerFusion, a::AbstractArray, biperm::AbstractBlockPermutation{2}
-)
-  return matricize(arg1(style), arg1(a), biperm) ⊗ matricize(arg2(style), arg2(a), biperm)
-end
-function TensorAlgebra.matricize(
-  style::KroneckerFusion, a::AbstractArray, biperm::AbstractBlockPermutation{2}
-)
-  return matricize_kronecker(style, a, biperm)
-end
-# Fix ambiguity error.
-# TODO: Investigate rewriting the logic in `TensorAlgebra.jl` to avoid this.
-using TensorAlgebra: BlockedTrivialPermutation, unmatricize
-function TensorAlgebra.matricize(
-  style::KroneckerFusion, a::AbstractArray, biperm::BlockedTrivialPermutation{2}
-)
-  return matricize_kronecker(style, a, biperm)
-end
-function unmatricize_kronecker(style::KroneckerFusion, a::AbstractArray, ax)
-  return unmatricize(arg1(style), arg1(a), arg1.(ax)) ⊗
-         unmatricize(arg2(style), arg2(a), arg2.(ax))
-end
-function TensorAlgebra.unmatricize(style::KroneckerFusion, a::AbstractArray, ax)
-  return unmatricize_kronecker(style, a, ax)
-end