[ITensors] Make tr preserve element type of input (#1448)

ITensor · May 17, 2024 · ce24a30 · ce24a30
1 parent c8168a4
commit ce24a30
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/src/tensor_operations/matrix_algebra.jl b/src/tensor_operations/matrix_algebra.jl
@@ -6,7 +6,7 @@ function _tr(T::ITensor; plev::Pair{Int,Int}=0 => 1, tags::Pair=ts"" => ts"")
   Tᶜ = T * Cᴸ * Cᴿ
   cᴸ = uniqueind(Cᴸ, T)
   cᴿ = uniqueind(Cᴿ, T)
-  Tᶜ *= δ(dag((cᴸ, cᴿ)))
+  Tᶜ *= δ(eltype(T), dag((cᴸ, cᴿ)))
   if order(Tᶜ) == 0
     return Tᶜ[]
   end

diff --git a/test/base/test_itensor.jl b/test/base/test_itensor.jl
@@ -316,16 +316,20 @@ end
       @test ndims(A) == 0
     end
 
-    @testset "trace (tr)" begin
+    @testset "trace (tr) (eltype=$elt)" for elt in (
+      Float32, Float64, Complex{Float32}, Complex{Float64}
+    )
       i, j, k, l = Index.((2, 3, 4, 5), ("i", "j", "k", "l"))
-      T = random_itensor(j, k', i', k, j', i)
+      T = random_itensor(elt, j, k', i', k, j', i)
       trT1 = tr(T)
-      trT2 = (T * δ(i, i') * δ(j, j') * δ(k, k'))[]
+      @test eltype(trT1) === elt
+      trT2 = (T * δ(elt, i, i') * δ(elt, j, j') * δ(elt, k, k'))[]
       @test trT1 ≈ trT2
 
-      T = random_itensor(j, k', i', l, k, j', i)
+      T = random_itensor(elt, j, k', i', l, k, j', i)
       trT1 = tr(T)
-      trT2 = T * δ(i, i') * δ(j, j') * δ(k, k')
+      @test eltype(trT1) === elt
+      trT2 = T * δ(elt, i, i') * δ(elt, j, j') * δ(elt, k, k')
       @test trT1 ≈ trT2
     end