Merge 6d57dce into 70eb783

ITensor · Jul 20, 2023 · 7507112 · 7507112
2 parents 70eb783 + 6d57dce
commit 7507112
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 3 deletions.
diff --git a/NDTensors/src/dense/tensoralgebra/contract.jl b/NDTensors/src/dense/tensoralgebra/contract.jl
@@ -375,9 +375,17 @@ function _contract!(
 
   # TODO: this logic may be wrong
   if props.permuteC
-    # Need to copy here since we will be permuting
-    # into C later
-    CM = reshape(copy(CT), (props.dleft, props.dright))
+    # if we are computing C = α * A B + β * C
+    # we need to make sure C is permuted to the same 
+    # ordering as A B
+    if β ≠ 0
+      pC = NTuple{NB,Int}(props.PC)
+      CM = reshape(permutedims(CT, pC), (props.dleft, props.dright))
+    else
+      # Need to copy here since we will be permuting
+      # into C later  
+      CM = reshape(copy(CT), (props.dleft, props.dright))
+    end
   else
     if Ctrans(props)
       CM = transpose(reshape(CT, (props.dright, props.dleft)))

diff --git a/test/base/test_contract.jl b/test/base/test_contract.jl
@@ -241,6 +241,16 @@ digits(::Type{T}, i, j, k) where {T} = T(i * 10^2 + j * 10 + k)
         @test CArray ≈ array(permute(C, α, i, j))
       end
     end
+    @testset "Test contract in-place ITensors (4-Tensor*Matrix -> 4-Tensor)" begin
+      A = randomITensor(T, (j, i))
+      B = randomITensor(T, (j, k, l, α))
+      C = ITensor(zero(T), (i, k, α, l))
+      ITensors.contract!(C, A, B, 1.0, 0.0)
+      ITensors.contract!(C, A, B, 1.0, 1.0)
+      D = A * B
+      D .+= A * B
+      @test C ≈ D
+    end
   end # End contraction testset
 end