diff --git a/NDTensors/src/dense/tensoralgebra/contract.jl b/NDTensors/src/dense/tensoralgebra/contract.jl index 15f9ac1b5e..31e601ced9 100644 --- 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 index 6078489f37..a68da19f1d 100644 --- 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