Small fixes for upstream + CUDA

Author: kshyatt

ext/TensorKitCUDAExt/cutensormap.jl

@@ -37,6 +37,11 @@ for (fname, felt) in ((:zeros, :zero), (:ones, :one))
             fill!(t, $felt(T))
             return t
         end
+        function Base.$fname(::Type{TorA}, V::TensorMapSpace) where {TorA <: CuArray}
+            t = tensormaptype(spacetype(V), numout(V), numin(V), TorA)(undef, V)
+            fill!(t, $felt(eltype(TorA)))
+            return t
+        end
     end
 end
 

ext/TensorKitChainRulesCoreExt/linalg.jl

@@ -80,7 +80,7 @@ function ChainRulesCore.rrule(
 end
 
 function ChainRulesCore.rrule(::typeof(tr), A::AbstractTensorMap)
-    tr_pullback(Δtr) = NoTangent(), Δtr * id(domain(A))
+    tr_pullback(Δtr) = NoTangent(), Δtr * id(storagetype(A), domain(A))
     return tr(A), tr_pullback
 end
 

src/tensors/tensor.jl

@@ -306,14 +306,19 @@ for (fname, felt) in ((:zeros, :zero), (:ones, :one))
             return Base.$fname(codomain ← domain)
         end
         function Base.$fname(
-                ::Type{T}, codomain::TensorSpace{S}, domain::TensorSpace{S} = one(codomain)
-            ) where {T, S <: IndexSpace}
-            return Base.$fname(T, codomain ← domain)
+                ::Type{TorA}, codomain::TensorSpace{S}, domain::TensorSpace{S} = one(codomain)
+            ) where {TorA, S <: IndexSpace}
+            return Base.$fname(TorA, codomain ← domain)
+        end
+        function Base.$fname(
+                ::Type{T}, ::Type{TorA}, codomain::TensorSpace{S}, domain::TensorSpace{S} = one(codomain)
+            ) where {T, TorA, S <: IndexSpace}
+            return Base.$fname(TorA, codomain ← domain)
         end
         Base.$fname(V::TensorMapSpace) = Base.$fname(Float64, V)
-        function Base.$fname(::Type{T}, V::TensorMapSpace) where {T}
-            t = TensorMap{T}(undef, V)
-            fill!(t, $felt(T))
+        function Base.$fname(::Type{TorA}, V::TensorMapSpace) where {TorA}
+            t = tensormaptype(spacetype(V), numout(V), numin(V), TorA)(undef, V)
+            fill!(t, $felt(TorA))
             return t
         end
     end