TensorOperations for TensorAlgebra

Author: lkdvos

ext/TensorAlgebraTensorOperationsExt.jl

@@ -19,6 +19,78 @@ end
 
 _blockedpermutation(p::Index2Tuple) = TensorAlgebra.blockedpermvcat(p...)
 
+# Using TensorOperations backends as TensorAlgebra implementations
+# ----------------------------------------------------------------
+
+# not in-place
+function TensorAlgebra.contract(
+  backend::TOAlgorithm,
+  bipermAB::BlockedPermutation,
+  A::AbstractArray,
+  bipermA::BlockedPermutation,
+  B::AbstractArray,
+  bipermB::BlockedPermutation,
+  α::Number,
+)
+  pA = _index2tuple(bipermA)
+  pB = _index2tuple(bipermB)
+
+  # TODO: this assumes biperm of output because not enough information!
+  ipermAB = invperm(Tuple(bipermAB))
+  pAB = (TupleTools.getindices(ipermAB, length(ipermAB)), ())
+
+  return tensorcontract(A, pA, false, B, pB, false, pAB, α, backend)
+end
+
+function TensorAlgebra.contract(
+  backend::TOAlgorithm,
+  labelsC,
+  A::AbstractArray,
+  labelsA,
+  B::AbstractArray,
+  labelsB,
+  α::Number,
+)
+  return tensorcontract(labelsC, A, labelsA, B, labelsB, α; backend)
+end
+
+# in-place
+function TensorAlgebra.contract!(
+  backend::TOAlgorithm,
+  C::AbstractArray,
+  bipermAB::BlockedPermutation,
+  A::AbstractArray,
+  bipermA::BlockedPermutation,
+  B::AbstractArray,
+  bipermB::BlockedPermutation,
+  α::Number,
+  β::Number,
+)
+  pA = _index2tuple(bipermA)
+  pB = _index2tuple(bipermB)
+
+  # TODO: this assumes biperm of output because not enough information!
+  ipermAB = invperm(Tuple(bipermAB))
+  pAB = (TupleTools.getindices(ipermAB, length(ipermAB)), ())
+
+  return tensorcontract!(C, A, pA, false, B, pB, false, pAB, α, β, backend)
+end
+
+function TensorAlgebra.contract!(
+  backend::TOAlgorithm,
+  C::AbstractArray,
+  labelsC,
+  A::AbstractArray,
+  labelsA,
+  B::AbstractArray,
+  labelsB,
+  α::Number,
+  β::Number,
+)
+  pA, pB, pAB = TensorOperations.contract_indices(labelsA, labelsB, labelsC)
+  return TensorOperations.tensorcontract!(C, A, pA, false, B, pB, false, pAB, α, β, backend)
+end
+
 # Using TensorAlgebra implementations as TensorOperations backends
 # ----------------------------------------------------------------
 function TensorOperations.tensorcontract!(
@@ -38,10 +110,13 @@ function TensorOperations.tensorcontract!(
   bipermA = _blockedpermutation(pA)
   bipermB = _blockedpermutation(pB)
   ipAB = invperm((pAB[1]..., pAB[2]...))
-  bipermAB = _blockedpermutation((TupleTools.getindices(ipAB, trivtuple(length(pA[1]))),
-  TupleTools.getindices(ipAB, trivtuple(length(pB[2])) .+ length(pA[1]))))
+  bipermAB = _blockedpermutation((
+    TupleTools.getindices(ipAB, trivtuple(length(pA[1]))),
+    TupleTools.getindices(ipAB, trivtuple(length(pB[2])) .+ length(pA[1])),
+  ))
   A′ = conjA ? conj(A) : A
   B′ = conjB ? conj(B) : B
   return TensorAlgebra.contract!(backend, C, bipermAB, A′, bipermA, B′, bipermB, α, β)
 end
 
+end