Skip to content

Commit c47e269

Browse files
kshyattkshyatt
committed
More comments
1 parent dd056b1 commit c47e269

File tree

4 files changed

+46
-45
lines changed

4 files changed

+46
-45
lines changed

‎ext/TensorKitCUDAExt/cutensormap.jl‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -142,7 +142,7 @@ end
142142
function TensorKit.exp!(t::CuTensorMap)
143143
domain(t) == codomain(t) ||
144144
error("Exponential of a tensor only exist when domain == codomain.")
145-
!MatrixAlgebraKit.ishermitian(t) && throw(ArgumentError("`exp!` is only supported on hermitian CUDA tensors"))
145+
!MatrixAlgebraKit.ishermitian(t) && throw(ArgumentError("`exp!` is currently only supported on hermitian CUDA tensors"))
146146
for (c, b) in blocks(t)
147147
copy!(b, parent(Base.exp(Hermitian(b))))
148148
end
@@ -155,7 +155,7 @@ for f in (:sqrt, :log, :asin, :acos, :acosh, :atanh, :acoth)
155155
@eval function Base.$f(t::CuTensorMap)
156156
domain(t) == codomain(t) ||
157157
throw(SpaceMismatch("`$($sf)` of a tensor only exists when domain == codomain"))
158-
!MatrixAlgebraKit.ishermitian(t) && throw(ArgumentError("`$($sf)` is only supported on hermitian CUDA tensors"))
158+
!MatrixAlgebraKit.ishermitian(t) && throw(ArgumentError("`$($sf)` is currently only supported on hermitian CUDA tensors"))
159159
T = complex(float(scalartype(t)))
160160
tf = similar(t, T)
161161
for (c, b) in blocks(t)

‎src/tensors/linalg.jl‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -289,7 +289,7 @@ function LinearAlgebra.rank(
289289
)
290290
r = 0 * dim(first(allunits(sectortype(t))))
291291
dim(t) == 0 && return r
292-
S = svd_vals(t)
292+
S = MatrixAlgebraKit.svd_vals(t)
293293
tol = max(atol, rtol * maximum(parent(S)))
294294
for (c, b) in pairs(S)
295295
if !isempty(b)
@@ -307,7 +307,7 @@ function LinearAlgebra.cond(t::AbstractTensorMap, p::Real = 2)
307307
throw(SpaceMismatch("`cond` requires domain and codomain to be the same"))
308308
return zero(real(float(scalartype(t))))
309309
end
310-
S = svd_vals(t)
310+
S = MatrixAlgebraKit.svd_vals(t)
311311
maxS = maximum(parent(S))
312312
minS = minimum(parent(S))
313313
return iszero(maxS) ? oftype(maxS, Inf) : (maxS / minS)

‎src/tensors/tensor.jl‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -404,9 +404,10 @@ for randf in (:rand, :randn, :randexp, :randisometry)
404404
end
405405
end
406406

407-
# Collecting arbitrary TensorMaps
407+
# Moving arbitrary TensorMaps to CPU
408408
#-----------------------------
409-
Base.collect(t::TensorMap) = convert(TensorMapWithStorage{scalartype(t), similarstoragetype(scalartype(t))}, t)
409+
to_cpu(t::TensorMapWithStorage{T, Vector{T}}) = t # no op
410+
to_cpu(t::TensorMap) = convert(TensorMapWithStorage{scalartype(t), similarstoragetype(scalartype(t))}, t)
410411

411412
# Efficient copy constructors
412413
#-----------------------------

‎test/cuda/tensors.jl‎

Lines changed: 39 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -120,7 +120,7 @@ for V in spacelist
120120
for T in (Int, Float32, ComplexF64)
121121
t = @constinferred CUDA.rand(T, W)
122122
d = convert(Dict, t)
123-
@test collect(t) == convert(TensorMap, d)
123+
@test TensorKit.to_cpu(t) == convert(TensorMap, d)
124124
end
125125
end
126126
@timedtestset "Basic linear algebra" begin
@@ -212,10 +212,10 @@ for V in spacelist
212212
t = CUDA.rand(T, W)
213213
t2 = @constinferred CUDA.rand!(similar(t))
214214
α = rand(T)
215-
@test norm(t, 2) ≈ norm(collect(t), 2)
216-
@test dot(t2, t) ≈ dot(collect(t2), collect(t))
217-
@test collect(α * t) ≈ α * collect(t)
218-
@test collect(t + t) ≈ 2 * collect(t)
215+
@test norm(t, 2) ≈ norm(TensorKit.to_cpu(t), 2)
216+
@test dot(t2, t) ≈ dot(TensorKit.to_cpu(t2), TensorKit.to_cpu(t))
217+
@test TensorKit.to_cpu(α * t) ≈ α * TensorKit.to_cpu(t)
218+
@test TensorKit.to_cpu(t + t) ≈ 2 * TensorKit.to_cpu(t)
219219
end
220220
end
221221
@timedtestset "Real and imaginary parts" begin
@@ -225,17 +225,17 @@ for V in spacelist
225225

226226
tr = @constinferred real(t)
227227
@test scalartype(tr) <: Real
228-
@test real(collect(t)) == collect(tr)
228+
@test real(TensorKit.to_cpu(t)) == TensorKit.to_cpu(tr)
229229
@test storagetype(tr) == CuVector{real(T), CUDA.DeviceMemory}
230230

231231
ti = @constinferred imag(t)
232232
@test scalartype(ti) <: Real
233-
@test imag(collect(t)) == collect(ti)
233+
@test imag(TensorKit.to_cpu(t)) == TensorKit.to_cpu(ti)
234234
@test storagetype(ti) == CuVector{real(T), CUDA.DeviceMemory}
235235

236236
tc = @inferred complex(t)
237237
@test scalartype(tc) <: Complex
238-
@test complex(collect(t)) == collect(tc)
238+
@test complex(TensorKit.to_cpu(t)) == TensorKit.to_cpu(tc)
239239
@test storagetype(tc) == CuVector{complex(T), CUDA.DeviceMemory}
240240

241241
tc2 = @inferred complex(tr, ti)
@@ -295,19 +295,19 @@ for V in spacelist
295295
@timedtestset "Permutations: test via CPU" begin
296296
W = V1 ⊗ V2 ⊗ V3 ⊗ V4 ⊗ V5
297297
t = CUDA.rand(ComplexF64, W)
298-
a = convert(Array, collect(t))
298+
a = convert(Array, TensorKit.to_cpu(t))
299299
for k in 0:5
300300
for p in permutations(1:5)
301301
p1 = ntuple(n -> p[n], k)
302302
p2 = ntuple(n -> p[k + n], 5 - k)
303303
dt2 = CUDA.@allowscalar permute(t, (p1, p2))
304-
ht2 = permute(collect(t), (p1, p2))
305-
@test ht2 == collect(dt2)
304+
ht2 = permute(TensorKit.to_cpu(t), (p1, p2))
305+
@test ht2 == TensorKit.to_cpu(dt2)
306306
end
307307

308308
dt3 = CUDA.@allowscalar repartition(t, k)
309-
ht3 = repartition(collect(t), k)
310-
@test ht3 == collect(dt3)
309+
ht3 = repartition(TensorKit.to_cpu(t), k)
310+
@test ht3 == TensorKit.to_cpu(dt3)
311311
end
312312
end
313313
end
@@ -368,10 +368,10 @@ for V in spacelist
368368
@tensor dHrA12[a, s1, s2, c] := drhoL[a, a'] * conj(dA1[a', t1, b]) *
369369
dA2[b, t2, c'] * drhoR[c', c] *
370370
dH[s1, s2, t1, t2]
371-
@tensor hHrA12[a, s1, s2, c] := collect(drhoL)[a, a'] * conj(collect(dA1)[a', t1, b]) *
372-
collect(dA2)[b, t2, c'] * collect(drhoR)[c', c] *
373-
collect(dH)[s1, s2, t1, t2]
374-
@test collect(dHrA12) ≈ hHrA12
371+
@tensor hHrA12[a, s1, s2, c] := TensorKit.to_cpu(drhoL)[a, a'] * conj(TensorKit.to_cpu(dA1)[a', t1, b]) *
372+
TensorKit.to_cpu(dA2)[b, t2, c'] * TensorKit.to_cpu(drhoR)[c', c] *
373+
TensorKit.to_cpu(dH)[s1, s2, t1, t2]
374+
@test TensorKit.to_cpu(dHrA12) ≈ hHrA12
375375
end
376376
end=# # doesn't yet work because of AdjointTensor
377377
@timedtestset "Index flipping: test flipping inverse" begin
@@ -450,48 +450,48 @@ for V in spacelist
450450
t1 = CUDA.rand(T, W1, W1)
451451
t2 = CUDA.rand(T, W2, W2)
452452
t = CUDA.rand(T, W1, W2)
453-
ht1 = collect(t1)
454-
ht2 = collect(t2)
455-
ht = collect(t)
456-
@test collect(t1 * t) ≈ ht1 * ht
457-
@test collect(t1' * t) ≈ ht1' * ht
458-
@test collect(t2 * t') ≈ ht2 * ht'
459-
@test collect(t2' * t') ≈ ht2' * ht'
453+
ht1 = TensorKit.to_cpu(t1)
454+
ht2 = TensorKit.to_cpu(t2)
455+
ht = TensorKit.to_cpu(t)
456+
@test TensorKit.to_cpu(t1 * t) ≈ ht1 * ht
457+
@test TensorKit.to_cpu(t1' * t) ≈ ht1' * ht
458+
@test TensorKit.to_cpu(t2 * t') ≈ ht2 * ht'
459+
@test TensorKit.to_cpu(t2' * t') ≈ ht2' * ht'
460460

461-
@test collect(inv(t1)) ≈ inv(ht1)
462-
@test collect(pinv(t)) ≈ pinv(ht)
461+
@test TensorKit.to_cpu(inv(t1)) ≈ inv(ht1)
462+
@test TensorKit.to_cpu(pinv(t)) ≈ pinv(ht)
463463

464464
if T == Float32 || T == ComplexF32
465465
continue
466466
end
467467

468-
@test collect(t1 \ t) ≈ ht1 \ ht
469-
@test collect(t1' \ t) ≈ ht1' \ ht
470-
@test collect(t2 \ t') ≈ ht2 \ ht'
471-
@test collect(t2' \ t') ≈ ht2' \ ht'
468+
@test TensorKit.to_cpu(t1 \ t) ≈ ht1 \ ht
469+
@test TensorKit.to_cpu(t1' \ t) ≈ ht1' \ ht
470+
@test TensorKit.to_cpu(t2 \ t') ≈ ht2 \ ht'
471+
@test TensorKit.to_cpu(t2' \ t') ≈ ht2' \ ht'
472472

473-
@test collect(t2 / t) ≈ ht2 / ht
474-
@test collect(t2' / t) ≈ ht2' / ht
475-
@test collect(t1 / t') ≈ ht1 / ht'
476-
@test collect(t1' / t') ≈ ht1' / ht'
473+
@test TensorKit.to_cpu(t2 / t) ≈ ht2 / ht
474+
@test TensorKit.to_cpu(t2' / t) ≈ ht2' / ht
475+
@test TensorKit.to_cpu(t1 / t') ≈ ht1 / ht'
476+
@test TensorKit.to_cpu(t1' / t') ≈ ht1' / ht'
477477
end
478478
end
479479
end
480480
if BraidingStyle(I) isa Bosonic && hasfusiontensor(I)
481-
@timedtestset "Tensor functions" begin
481+
#=@timedtestset "Tensor functions" begin
482482
W = V1 ⊗ V2
483483
for T in (Float64, ComplexF64)
484484
t = CUDA.randn(T, W, W)
485485
s = dim(W)
486486
@test_broken (@constinferred sqrt(t))^2 ≈ t
487-
@test_broken collect(sqrt(t)) ≈ sqrt(collect(t))
487+
@test_broken TensorKit.to_cpu(sqrt(t)) ≈ sqrt(TensorKit.to_cpu(t))
488488
489489
expt = @constinferred exp(t)
490-
@test_broken collect(expt) ≈ exp(collect(t))
490+
@test_broken TensorKit.to_cpu(expt) ≈ exp(TensorKit.to_cpu(t))
491491
492492
# log doesn't work on CUDA yet (scalar indexing)
493493
#@test exp(@constinferred log(expt)) ≈ expt
494-
#@test collect(log(expt)) ≈ log(collect(expt))
494+
#@test TensorKit.to_cpu(log(expt)) ≈ log(TensorKit.to_cpu(expt))
495495
496496
#=@test (@constinferred cos(t))^2 + (@constinferred sin(t))^2 ≈
497497
id(storagetype(t), W)
@@ -520,7 +520,7 @@ for V in spacelist
520520
@test coth(@constinferred acoth(t8)) ≈ t8=#
521521
# TODO in CUDA
522522
end
523-
end
523+
end=#
524524
end
525525
# Sylvester not defined for CUDA
526526
# @timedtestset "Sylvester equation" begin

0 commit comments

Comments
 (0)