Comments

Author: kshyatt

ext/TensorKitCUDAExt/TensorKitCUDAExt.jl

@@ -12,7 +12,7 @@ using TensorKit.Factorizations: AbstractAlgorithm
 using TensorKit: SectorDict, tensormaptype, scalar, similarstoragetype, AdjointTensorMap, scalartype, project_symmetric_and_check
 import TensorKit: randisometry, rand, randn
 
-using TensorKit.MatrixAlgebraKit
+using TensorKit: MatrixAlgebraKit
 
 using Random
 

ext/TensorKitCUDAExt/cutensormap.jl

@@ -84,6 +84,13 @@ for randfun in (:curand, :curandn)
             $randfun!(rng, t)
             return t
         end
+
+        function $randfun!(rng::Random.AbstractRNG, t::CuTensorMap)
+            for (_, b) in blocks(t)
+                $randfun!(rng, b)
+            end
+            return t
+        end
     end
 end
 
@@ -112,7 +119,7 @@ function LinearAlgebra.isposdef(t::CuTensorMap)
     InnerProductStyle(spacetype(t)) === EuclideanInnerProduct() || return false
     for (c, b) in blocks(t)
         # do our own hermitian check
-        isherm = TensorKit.MatrixAlgebraKit.ishermitian(b; atol = eps(real(eltype(b))), rtol = eps(real(eltype(b))))
+        isherm = MatrixAlgebraKit.ishermitian(b; atol = eps(real(eltype(b))), rtol = eps(real(eltype(b))))
         isherm || return false
         isposdef(Hermitian(b)) || return false
     end
@@ -135,6 +142,7 @@ end
 function TensorKit.exp!(t::CuTensorMap)
     domain(t) == codomain(t) ||
         error("Exponential of a tensor only exist when domain == codomain.")
+    !MatrixAlgebraKit.ishermitian(t) && throw(ArgumentError("`exp!` is only supported on hermitian CUDA tensors"))
     for (c, b) in blocks(t)
         copy!(b, parent(Base.exp(Hermitian(b))))
     end
@@ -146,7 +154,8 @@ for f in (:sqrt, :log, :asin, :acos, :acosh, :atanh, :acoth)
     sf = string(f)
     @eval function Base.$f(t::CuTensorMap)
         domain(t) == codomain(t) ||
-            throw(SpaceMismatch("`$($sf)` of a tensor only exist when domain == codomain"))
+            throw(SpaceMismatch("`$($sf)` of a tensor only exists when domain == codomain"))
+        !MatrixAlgebraKit.ishermitian(t) && throw(ArgumentError("`$($sf)` is only supported on hermitian CUDA tensors"))
         T = complex(float(scalartype(t)))
         tf = similar(t, T)
         for (c, b) in blocks(t)

src/tensors/linalg.jl

@@ -289,7 +289,7 @@ function LinearAlgebra.rank(
     )
     r = 0 * dim(first(allunits(sectortype(t))))
     dim(t) == 0 && return r
-    S = LinearAlgebra.svdvals(t)
+    S = svd_vals(t)
     tol = max(atol, rtol * maximum(parent(S)))
     for (c, b) in pairs(S)
         if !isempty(b)
@@ -307,7 +307,7 @@ function LinearAlgebra.cond(t::AbstractTensorMap, p::Real = 2)
                 throw(SpaceMismatch("`cond` requires domain and codomain to be the same"))
             return zero(real(float(scalartype(t))))
         end
-        S = LinearAlgebra.svdvals(t)
+        S = svd_vals(t)
         maxS = maximum(parent(S))
         minS = minimum(parent(S))
         return iszero(maxS) ? oftype(maxS, Inf) : (maxS / minS)

test/cuda/tensors.jl

@@ -7,6 +7,8 @@ const CUDAExt = Base.get_extension(TensorKit, :TensorKitCUDAExt)
 const CuTensorMap = getglobal(CUDAExt, :CuTensorMap)
 const curand = getglobal(CUDAExt, :curand)
 const curandn = getglobal(CUDAExt, :curandn)
+const curand! = getglobal(CUDAExt, :curand!)
+const curandn! = getglobal(CUDAExt, :curandn!)
 
 @isdefined(TestSetup) || include("../setup.jl")
 using .TestSetup
@@ -61,6 +63,15 @@ for V in spacelist
                 @test domain(t) == one(W)
                 @test typeof(t) == TensorMap{Float64, spacetype(t), 5, 0, CuVector{Float64, CUDA.DeviceMemory}}
             end
+            for f! in (curand!, curandn!)
+                t = @constinferred CUDA.zeros(W)
+                f!(t)
+                @test scalartype(t) == Float64
+                @test codomain(t) == W
+                @test space(t) == (W ← one(W))
+                @test domain(t) == one(W)
+                @test typeof(t) == TensorMap{Float64, spacetype(t), 5, 0, CuVector{Float64, CUDA.DeviceMemory}}
+            end
             for T in (Int, Float32, Float64, ComplexF32, ComplexF64)
                 t = @constinferred CUDA.zeros(T, W)
                 CUDA.@allowscalar begin