Small fixes for upstream + CUDA

Author: kshyatt

ext/TensorKitCUDAExt/TensorKitCUDAExt.jl

@@ -10,7 +10,7 @@ using TensorKit.Factorizations
 using TensorKit.Strided
 using TensorKit.Factorizations: AbstractAlgorithm
 using TensorKit: SectorDict, tensormaptype, scalar, similarstoragetype, AdjointTensorMap, scalartype, project_symmetric_and_check
-import TensorKit: randisometry, rand, randn
+import TensorKit: randisometry, rand, randn, twist!
 
 using TensorKit: MatrixAlgebraKit
 

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/constructors.jl

@@ -6,7 +6,7 @@
 
 function ChainRulesCore.rrule(::Type{TensorMap}, d::DenseArray, args...; kwargs...)
     function TensorMap_pullback(Δt)
-        ∂d = convert(Array, unthunk(Δt))
+        ∂d = convert(typeof(d), unthunk(Δt))
         return NoTangent(), ∂d, ntuple(_ -> NoTangent(), length(args))...
     end
     return TensorMap(d, args...; kwargs...), TensorMap_pullback

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