Re-enable Strided Diagonal implementation

Jutho · Apr 21, 2024 · 959bc09 · 959bc09
1 parent 52910cc
commit 959bc09
Showing 1 changed file with 14 additions and 4 deletions.
diff --git a/src/implementation/diagonal.jl b/src/implementation/diagonal.jl
@@ -1,6 +1,20 @@
 #-------------------------------------------------------------------------------------------
 # Specialized implementations for contractions involving diagonal matrices
 #-------------------------------------------------------------------------------------------
+
+# backend selection:
+for (TC, TA, TB) in ((:AbstractArray, :AbstractArray, :Diagonal),
+                     (:AbstractArray, :Diagonal, :AbstractArray), (:AbstractArray, :Diagonal, :Diagonal),
+                     (:Diagonal, :Diagonal, :Diagonal))
+    @eval function tensorcontract!(C::$TC, pC::Index2Tuple,
+                                   A::$TA, pA::Index2Tuple, conjA::Symbol,
+                                   B::$TB, pB::Index2Tuple, conjB::Symbol,
+                                   α::Number, β::Number)
+        return tensorcontract!(C, pC, A, pA, conjA, B, pB, conjB, α, β, StridedNative())
+    end
+end
+
+# actual implementations:
 function tensorcontract!(C::AbstractArray, pC::Index2Tuple,
                          A::AbstractArray, pA::Index2Tuple, conjA::Symbol,
                          B::Diagonal, pB::Index2Tuple, conjB::Symbol,
@@ -12,7 +26,6 @@ function tensorcontract!(C::AbstractArray, pC::Index2Tuple,
                          StridedView(B.diag), pB, conjB, α, β)
     return C
 end
-
 function tensorcontract!(C::AbstractArray, pC::Index2Tuple,
                          A::Diagonal, pA::Index2Tuple, conjA::Symbol,
                          B::AbstractArray, pB::Index2Tuple, conjB::Symbol,
@@ -33,7 +46,6 @@ function tensorcontract!(C::AbstractArray, pC::Index2Tuple,
                          StridedView(A.diag), rpA, conjA, α, β)
     return C
 end
-
 function tensorcontract!(C::AbstractArray, pC::Index2Tuple,
                          A::Diagonal, pA::Index2Tuple, conjA::Symbol,
                          B::Diagonal, pB::Index2Tuple, conjB::Symbol,
@@ -76,7 +88,6 @@ function tensorcontract!(C::AbstractArray, pC::Index2Tuple,
 
     return C
 end
-
 function tensorcontract!(C::Diagonal, pC::Index2Tuple,
                          A::Diagonal, pA::Index2Tuple, conjA::Symbol,
                          B::Diagonal, pB::Index2Tuple, conjB::Symbol,
@@ -91,7 +102,6 @@ function tensorcontract!(C::Diagonal, pC::Index2Tuple,
     C2 .= C2 .* β .+ A2 .* B2 .* α
     return C
 end
-
 function _diagtensorcontract!(C::StridedView, pC::Index2Tuple,
                               A::StridedView, pA::Index2Tuple, conjA::Symbol,
                               Bdiag::StridedView, pB::Index2Tuple, conjB::Symbol,