From 3e7795133c21b11cdd091e904b719ef3c5d170ed Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 13 Jun 2024 07:38:48 -0400 Subject: [PATCH 01/41] Working on GPU diag problem --- .../NDTensorsGPUArraysCoreExt.jl | 1 + .../ext/NDTensorsGPUArraysCoreExt/permutedims.jl | 15 +++++++++++++++ NDTensors/src/diag/diagtensor.jl | 2 +- 3 files changed, 17 insertions(+), 1 deletion(-) create mode 100644 NDTensors/ext/NDTensorsGPUArraysCoreExt/permutedims.jl diff --git a/NDTensors/ext/NDTensorsGPUArraysCoreExt/NDTensorsGPUArraysCoreExt.jl b/NDTensors/ext/NDTensorsGPUArraysCoreExt/NDTensorsGPUArraysCoreExt.jl index 1ead107b45..3511a5a9d7 100644 --- a/NDTensors/ext/NDTensorsGPUArraysCoreExt/NDTensorsGPUArraysCoreExt.jl +++ b/NDTensors/ext/NDTensorsGPUArraysCoreExt/NDTensorsGPUArraysCoreExt.jl @@ -1,3 +1,4 @@ module NDTensorsGPUArraysCoreExt include("contract.jl") +include("permutedims.jl") end diff --git a/NDTensors/ext/NDTensorsGPUArraysCoreExt/permutedims.jl b/NDTensors/ext/NDTensorsGPUArraysCoreExt/permutedims.jl new file mode 100644 index 0000000000..61488ca235 --- /dev/null +++ b/NDTensors/ext/NDTensorsGPUArraysCoreExt/permutedims.jl @@ -0,0 +1,15 @@ +using NDTensors: NDTensors, array +using GPUArraysCore: @allowscalar, AbstractGPUArray +using NDTensors.Adapt +function NDTensors.permutedims!( + Rexposed::Exposed{<:AbstractGPUArray, <:DiagTensor}, + texposed::Exposed{<:AbstractGPUArray, <:DiagTensor}, + perm::NTuple{N,Int}, + f::Function=(r, t) -> t, +) where {N} + R = unexpose(Rexposed) + t = unexpose(texposed) + + array(R) .= f.(array(R), array(t)) + return R +end \ No newline at end of file diff --git a/NDTensors/src/diag/diagtensor.jl b/NDTensors/src/diag/diagtensor.jl index 7a5cddff0e..d5555ab5f8 100644 --- a/NDTensors/src/diag/diagtensor.jl +++ b/NDTensors/src/diag/diagtensor.jl @@ -175,7 +175,7 @@ function permutedims!!( f::Function=(r, t) -> t, ) where {N} R = convert(promote_type(typeof(R), typeof(T)), R) - permutedims!(R, T, perm, f) + permutedims!(expose(R), expose(T), perm, f) return R end From 721d0f00eb6d79e1d9db91a8a1b8da49d6e7276c Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 13 Jun 2024 07:39:03 -0400 Subject: [PATCH 02/41] format --- NDTensors/ext/NDTensorsGPUArraysCoreExt/permutedims.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/NDTensors/ext/NDTensorsGPUArraysCoreExt/permutedims.jl b/NDTensors/ext/NDTensorsGPUArraysCoreExt/permutedims.jl index 61488ca235..6565262622 100644 --- a/NDTensors/ext/NDTensorsGPUArraysCoreExt/permutedims.jl +++ b/NDTensors/ext/NDTensorsGPUArraysCoreExt/permutedims.jl @@ -2,8 +2,8 @@ using NDTensors: NDTensors, array using GPUArraysCore: @allowscalar, AbstractGPUArray using NDTensors.Adapt function NDTensors.permutedims!( - Rexposed::Exposed{<:AbstractGPUArray, <:DiagTensor}, - texposed::Exposed{<:AbstractGPUArray, <:DiagTensor}, + Rexposed::Exposed{<:AbstractGPUArray,<:DiagTensor}, + texposed::Exposed{<:AbstractGPUArray,<:DiagTensor}, perm::NTuple{N,Int}, f::Function=(r, t) -> t, ) where {N} @@ -12,4 +12,4 @@ function NDTensors.permutedims!( array(R) .= f.(array(R), array(t)) return R -end \ No newline at end of file +end From 9e222edcebf180e555dfd4219d346078a80a2506 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 13 Jun 2024 07:46:34 -0400 Subject: [PATCH 03/41] Use data over array --- NDTensors/ext/NDTensorsGPUArraysCoreExt/permutedims.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/NDTensors/ext/NDTensorsGPUArraysCoreExt/permutedims.jl b/NDTensors/ext/NDTensorsGPUArraysCoreExt/permutedims.jl index 6565262622..c1c5d3f1a3 100644 --- a/NDTensors/ext/NDTensorsGPUArraysCoreExt/permutedims.jl +++ b/NDTensors/ext/NDTensorsGPUArraysCoreExt/permutedims.jl @@ -1,4 +1,4 @@ -using NDTensors: NDTensors, array +using NDTensors: NDTensors, data using GPUArraysCore: @allowscalar, AbstractGPUArray using NDTensors.Adapt function NDTensors.permutedims!( @@ -10,6 +10,6 @@ function NDTensors.permutedims!( R = unexpose(Rexposed) t = unexpose(texposed) - array(R) .= f.(array(R), array(t)) + data(R) .= f.(data(R), data(t)) return R end From 7e95335f686edcaaedcea96a2ffee389a1ff9f7a Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 13 Jun 2024 07:48:52 -0400 Subject: [PATCH 04/41] remove adapt --- NDTensors/ext/NDTensorsGPUArraysCoreExt/permutedims.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/NDTensors/ext/NDTensorsGPUArraysCoreExt/permutedims.jl b/NDTensors/ext/NDTensorsGPUArraysCoreExt/permutedims.jl index c1c5d3f1a3..03ba18103f 100644 --- a/NDTensors/ext/NDTensorsGPUArraysCoreExt/permutedims.jl +++ b/NDTensors/ext/NDTensorsGPUArraysCoreExt/permutedims.jl @@ -1,6 +1,5 @@ using NDTensors: NDTensors, data using GPUArraysCore: @allowscalar, AbstractGPUArray -using NDTensors.Adapt function NDTensors.permutedims!( Rexposed::Exposed{<:AbstractGPUArray,<:DiagTensor}, texposed::Exposed{<:AbstractGPUArray,<:DiagTensor}, From aef7db6318afaf15e88536f0e68715622df53471 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 13 Jun 2024 07:51:01 -0400 Subject: [PATCH 05/41] Make expose permutedims for diagtensor to fix cpu error --- NDTensors/src/diag/diagtensor.jl | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/NDTensors/src/diag/diagtensor.jl b/NDTensors/src/diag/diagtensor.jl index d5555ab5f8..aa50b37be9 100644 --- a/NDTensors/src/diag/diagtensor.jl +++ b/NDTensors/src/diag/diagtensor.jl @@ -190,6 +190,15 @@ function permutedims!!( return R end +function NDTensors.permutedims!( + Rexposed::Exposed{<:Array, <:DiagTensor}, + texposed::Exposed{<:Array, <:DiagTensor}, + perm::NTuple{N,Int}, + f::Function=(r, t) -> t, +) where {N} + permutedims!(unexpose(Rexposed), unexpose(texposed), perm, f) +end + function permutedims!( R::DenseTensor{ElR,N}, T::DiagTensor{ElT,N}, perm::NTuple{N,Int}, f::Function=(r, t) -> t ) where {ElR,ElT,N} From f8a80c975604e145a2bf19e8e25f6241078b0695 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 13 Jun 2024 07:51:17 -0400 Subject: [PATCH 06/41] format --- NDTensors/src/diag/diagtensor.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/NDTensors/src/diag/diagtensor.jl b/NDTensors/src/diag/diagtensor.jl index aa50b37be9..2258cdfb8e 100644 --- a/NDTensors/src/diag/diagtensor.jl +++ b/NDTensors/src/diag/diagtensor.jl @@ -191,12 +191,12 @@ function permutedims!!( end function NDTensors.permutedims!( - Rexposed::Exposed{<:Array, <:DiagTensor}, - texposed::Exposed{<:Array, <:DiagTensor}, + Rexposed::Exposed{<:Array,<:DiagTensor}, + texposed::Exposed{<:Array,<:DiagTensor}, perm::NTuple{N,Int}, f::Function=(r, t) -> t, ) where {N} - permutedims!(unexpose(Rexposed), unexpose(texposed), perm, f) + return permutedims!(unexpose(Rexposed), unexpose(texposed), perm, f) end function permutedims!( From f96f3224d1936553d4ee0e3726f00dc1be805eba Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 13 Jun 2024 16:36:43 -0400 Subject: [PATCH 07/41] Remove permutedims function --- .../NDTensorsGPUArraysCoreExt.jl | 1 - .../ext/NDTensorsGPUArraysCoreExt/permutedims.jl | 14 -------------- 2 files changed, 15 deletions(-) delete mode 100644 NDTensors/ext/NDTensorsGPUArraysCoreExt/permutedims.jl diff --git a/NDTensors/ext/NDTensorsGPUArraysCoreExt/NDTensorsGPUArraysCoreExt.jl b/NDTensors/ext/NDTensorsGPUArraysCoreExt/NDTensorsGPUArraysCoreExt.jl index 3511a5a9d7..1ead107b45 100644 --- a/NDTensors/ext/NDTensorsGPUArraysCoreExt/NDTensorsGPUArraysCoreExt.jl +++ b/NDTensors/ext/NDTensorsGPUArraysCoreExt/NDTensorsGPUArraysCoreExt.jl @@ -1,4 +1,3 @@ module NDTensorsGPUArraysCoreExt include("contract.jl") -include("permutedims.jl") end diff --git a/NDTensors/ext/NDTensorsGPUArraysCoreExt/permutedims.jl b/NDTensors/ext/NDTensorsGPUArraysCoreExt/permutedims.jl deleted file mode 100644 index 03ba18103f..0000000000 --- a/NDTensors/ext/NDTensorsGPUArraysCoreExt/permutedims.jl +++ /dev/null @@ -1,14 +0,0 @@ -using NDTensors: NDTensors, data -using GPUArraysCore: @allowscalar, AbstractGPUArray -function NDTensors.permutedims!( - Rexposed::Exposed{<:AbstractGPUArray,<:DiagTensor}, - texposed::Exposed{<:AbstractGPUArray,<:DiagTensor}, - perm::NTuple{N,Int}, - f::Function=(r, t) -> t, -) where {N} - R = unexpose(Rexposed) - t = unexpose(texposed) - - data(R) .= f.(data(R), data(t)) - return R -end From 7807882a0ec00a1b4e2ad328b6b954c31cd834a7 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 13 Jun 2024 17:07:25 -0400 Subject: [PATCH 08/41] Try to make GPUs more supported by Diag --- NDTensors/src/diag/diagtensor.jl | 29 ++++++++--------------------- NDTensors/test/test_diag.jl | 14 +++++++++++--- 2 files changed, 19 insertions(+), 24 deletions(-) diff --git a/NDTensors/src/diag/diagtensor.jl b/NDTensors/src/diag/diagtensor.jl index 2258cdfb8e..3f084265b2 100644 --- a/NDTensors/src/diag/diagtensor.jl +++ b/NDTensors/src/diag/diagtensor.jl @@ -1,4 +1,4 @@ -using .DiagonalArrays: diaglength +using .DiagonalArrays: diaglength, diagview const DiagTensor{ElT,N,StoreT,IndsT} = Tensor{ElT,N,StoreT,IndsT} where {StoreT<:Diag} const NonuniformDiagTensor{ElT,N,StoreT,IndsT} = @@ -9,9 +9,7 @@ const UniformDiagTensor{ElT,N,StoreT,IndsT} = function diag(tensor::DiagTensor) tensor_diag = NDTensors.similar(dense(typeof(tensor)), (diaglength(tensor),)) # TODO: Define `eachdiagindex`. - for j in 1:diaglength(tensor) - tensor_diag[j] = getdiagindex(tensor, j) - end + diagview(array(tensor_diag)) .= diagview(array(tensor)) return tensor_diag end @@ -145,16 +143,14 @@ function permutedims!( f::Function=(r, t) -> t, ) where {N} # TODO: check that inds(R)==permute(inds(T),perm)? - for i in 1:diaglength(R) - @inbounds setdiagindex!(R, f(getdiagindex(R, i), getdiagindex(T, i)), i) - end + diagview(data(R)) .= f.(diagview(data(R)), diagview(data(T))) return R end function permutedims( T::DiagTensor{<:Number,N}, perm::NTuple{N,Int}, f::Function=identity ) where {N} - R = NDTensors.similar(T, permute(inds(T), perm)) + R = NDTensors.similar(T) g(r, t) = f(t) permutedims!(R, T, perm, g) return R @@ -175,7 +171,7 @@ function permutedims!!( f::Function=(r, t) -> t, ) where {N} R = convert(promote_type(typeof(R), typeof(T)), R) - permutedims!(expose(R), expose(T), perm, f) + permutedims!(R, T, perm, f) return R end @@ -190,21 +186,12 @@ function permutedims!!( return R end -function NDTensors.permutedims!( - Rexposed::Exposed{<:Array,<:DiagTensor}, - texposed::Exposed{<:Array,<:DiagTensor}, - perm::NTuple{N,Int}, - f::Function=(r, t) -> t, -) where {N} - return permutedims!(unexpose(Rexposed), unexpose(texposed), perm, f) -end - function permutedims!( R::DenseTensor{ElR,N}, T::DiagTensor{ElT,N}, perm::NTuple{N,Int}, f::Function=(r, t) -> t ) where {ElR,ElT,N} - for i in 1:diaglength(T) - @inbounds setdiagindex!(R, f(getdiagindex(R, i), getdiagindex(T, i)), i) - end + rview = diagview(array(R)) + tview = diagview(T) + rview .= f.(rview, tview) return R end diff --git a/NDTensors/test/test_diag.jl b/NDTensors/test/test_diag.jl index f6390a733f..4261c4f4b0 100644 --- a/NDTensors/test/test_diag.jl +++ b/NDTensors/test/test_diag.jl @@ -31,22 +31,30 @@ using LinearAlgebra: dot @test complex(D) == Diag(one(complex(elt))) @test similar(D) == Diag(0.0) - D = Tensor(Diag(1), (2, 2)) + D = dev(Tensor(Diag(1), (2, 2))) @test norm(D) == √2 d = 3 vr = rand(elt, d) D = dev(tensor(Diag(vr), (d, d))) Da = Array(D) Dm = Matrix(D) + Da = permutedims(D, (2, 1)) @allowscalar begin @test Da == NDTensors.LinearAlgebra.diagm(0 => vr) @test Da == NDTensors.LinearAlgebra.diagm(0 => vr) - ## TODO Currently this permutedims requires scalar indexing on GPU. - Da = permutedims(D, (2, 1)) @test Da == D end + if (dev == NDTensors.mtl && elt != ComplexF32) + S = permutedims(dev(D), (1, 2), sqrt) + @allowscalar begin + for i in 1:diaglength(S) + @test S[i, i] == sqrt(D[i, i]) + end + end + end + # Regression test for https://github.com/ITensor/ITensors.jl/issues/1199 S = dev(tensor(Diag(randn(elt, 2)), (2, 2))) ## This was creating a `Dense{ReshapedArray{Adjoint{Matrix}}}` which, in mul!, was From 6ef078ff8830415a634a77559783136ecb3e9601 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 13 Jun 2024 18:26:57 -0400 Subject: [PATCH 09/41] Add a comment with a link to the bug in Metal.jl --- NDTensors/test/test_diag.jl | 2 ++ 1 file changed, 2 insertions(+) diff --git a/NDTensors/test/test_diag.jl b/NDTensors/test/test_diag.jl index 4261c4f4b0..912aa39c8f 100644 --- a/NDTensors/test/test_diag.jl +++ b/NDTensors/test/test_diag.jl @@ -46,6 +46,8 @@ using LinearAlgebra: dot @test Da == D end + # This if statement corresponds to the reported bug: + # https://github.com/JuliaGPU/Metal.jl/issues/364 if (dev == NDTensors.mtl && elt != ComplexF32) S = permutedims(dev(D), (1, 2), sqrt) @allowscalar begin From c0ad602c7a7229e5b5d6aebdfd56bc189cbe17ca Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 13 Jun 2024 18:52:07 -0400 Subject: [PATCH 10/41] Remove unused line (A request from Miles) --- NDTensors/src/linearalgebra/linearalgebra.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/NDTensors/src/linearalgebra/linearalgebra.jl b/NDTensors/src/linearalgebra/linearalgebra.jl index bd1690d127..f038659475 100644 --- a/NDTensors/src/linearalgebra/linearalgebra.jl +++ b/NDTensors/src/linearalgebra/linearalgebra.jl @@ -369,7 +369,6 @@ matrix is unique. Returns a tuple (Q,R). function qr_positive(M::AbstractMatrix) sparseQ, R = qr(M) Q = convert(typeof(R), sparseQ) - nc = size(Q, 2) signs = nonzero_sign.(diag(R)) Q = Q * Diagonal(signs) R = Diagonal(conj.(signs)) * R From b3892df8e33b7fe3f3d265ff607ed1a3ac81a142 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 13 Jun 2024 19:33:55 -0400 Subject: [PATCH 11/41] Fix diag function for GPU code --- NDTensors/src/tensor/tensor.jl | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/NDTensors/src/tensor/tensor.jl b/NDTensors/src/tensor/tensor.jl index 4dfa483030..16e90ad14b 100644 --- a/NDTensors/src/tensor/tensor.jl +++ b/NDTensors/src/tensor/tensor.jl @@ -361,6 +361,7 @@ function getdiagindex(T::Tensor{<:Number,N}, ind::Int) where {N} return getindex(T, CartesianIndex(ntuple(_ -> ind, Val(N)))) end +using .DiagonalArrays: diagview # TODO: add support for off-diagonals, return # block sparse vector instead of dense. function diag(tensor::Tensor) @@ -368,9 +369,7 @@ function diag(tensor::Tensor) tensordiag = NDTensors.similar( dense(typeof(tensor)), eltype(tensor), (diaglength(tensor),) ) - for n in 1:diaglength(tensor) - tensordiag[n] = tensor[n, n] - end + data(tensordiag) .= diagview(tensor) return tensordiag end From cc24e62fb4a42866e6370dbbadcf0cfbab7b9cd4 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 13 Jun 2024 19:34:35 -0400 Subject: [PATCH 12/41] Make diagview functions for Tensor types --- NDTensors/src/diag/diagtensor.jl | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/NDTensors/src/diag/diagtensor.jl b/NDTensors/src/diag/diagtensor.jl index 3f084265b2..acb27174a5 100644 --- a/NDTensors/src/diag/diagtensor.jl +++ b/NDTensors/src/diag/diagtensor.jl @@ -9,7 +9,7 @@ const UniformDiagTensor{ElT,N,StoreT,IndsT} = function diag(tensor::DiagTensor) tensor_diag = NDTensors.similar(dense(typeof(tensor)), (diaglength(tensor),)) # TODO: Define `eachdiagindex`. - diagview(array(tensor_diag)) .= diagview(array(tensor)) + diagview(tensor_diag) .= diagview(tensor) return tensor_diag end @@ -31,6 +31,19 @@ function Array(T::DiagTensor{ElT,N}) where {ElT,N} return Array{ElT,N}(T) end +function DiagonalArrays.diagview(T::NonuniformDiagTensor) + return data(T) +end + +function DiagonalArrays.diagview(T::UniformDiagTensor) + return array(T) +end + +## Should this go in dense.jl or here since its related to diag? +function DiagonalArrays.diagview(T::DenseTensor) + return diagview(array(T)) +end + function zeros(tensortype::Type{<:DiagTensor}, inds) return tensor(generic_zeros(storagetype(tensortype), mindim(inds)), inds) end @@ -143,7 +156,7 @@ function permutedims!( f::Function=(r, t) -> t, ) where {N} # TODO: check that inds(R)==permute(inds(T),perm)? - diagview(data(R)) .= f.(diagview(data(R)), diagview(data(T))) + data(R) .= f.(diagview(R), diagview(T)) return R end @@ -190,8 +203,7 @@ function permutedims!( R::DenseTensor{ElR,N}, T::DiagTensor{ElT,N}, perm::NTuple{N,Int}, f::Function=(r, t) -> t ) where {ElR,ElT,N} rview = diagview(array(R)) - tview = diagview(T) - rview .= f.(rview, tview) + rview .= f.(rview, diagview(T)) return R end From 2af30bfa143e0f4b4cd48196069ac35059bdb621 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 13 Jun 2024 19:37:08 -0400 Subject: [PATCH 13/41] Remove unecessary function --- NDTensors/test/test_diag.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NDTensors/test/test_diag.jl b/NDTensors/test/test_diag.jl index 912aa39c8f..d4f6cbf188 100644 --- a/NDTensors/test/test_diag.jl +++ b/NDTensors/test/test_diag.jl @@ -31,7 +31,7 @@ using LinearAlgebra: dot @test complex(D) == Diag(one(complex(elt))) @test similar(D) == Diag(0.0) - D = dev(Tensor(Diag(1), (2, 2))) + D = Tensor(Diag(1), (2, 2)) @test norm(D) == √2 d = 3 vr = rand(elt, d) From 9f555bcba2618d4faef456047e702e6b025f4dd0 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 13 Jun 2024 19:51:30 -0400 Subject: [PATCH 14/41] Update permutedim functions --- NDTensors/src/diag/diagtensor.jl | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/NDTensors/src/diag/diagtensor.jl b/NDTensors/src/diag/diagtensor.jl index acb27174a5..d8f8e251ae 100644 --- a/NDTensors/src/diag/diagtensor.jl +++ b/NDTensors/src/diag/diagtensor.jl @@ -156,7 +156,7 @@ function permutedims!( f::Function=(r, t) -> t, ) where {N} # TODO: check that inds(R)==permute(inds(T),perm)? - data(R) .= f.(diagview(R), diagview(T)) + diagview(R) .= f.(diagview(R), diagview(T)) return R end @@ -202,8 +202,7 @@ end function permutedims!( R::DenseTensor{ElR,N}, T::DiagTensor{ElT,N}, perm::NTuple{N,Int}, f::Function=(r, t) -> t ) where {ElR,ElT,N} - rview = diagview(array(R)) - rview .= f.(rview, diagview(T)) + diagview(array(R)) .= f.(diagview(array(R)), diagview(T)) return R end From a3c13c609cad303a9968a11fd72deb4ef43ad6e0 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 13 Jun 2024 19:53:56 -0400 Subject: [PATCH 15/41] return fill for diagview uniformdiagtensor --- NDTensors/src/diag/diagtensor.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NDTensors/src/diag/diagtensor.jl b/NDTensors/src/diag/diagtensor.jl index d8f8e251ae..b2534b53d1 100644 --- a/NDTensors/src/diag/diagtensor.jl +++ b/NDTensors/src/diag/diagtensor.jl @@ -36,7 +36,7 @@ function DiagonalArrays.diagview(T::NonuniformDiagTensor) end function DiagonalArrays.diagview(T::UniformDiagTensor) - return array(T) + return fill(getdiagindex(T, 1), diaglength(T)) end ## Should this go in dense.jl or here since its related to diag? From 41bf2af7d866101db392d5ce6a094b9c3f74ac80 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 13 Jun 2024 19:55:10 -0400 Subject: [PATCH 16/41] remove uniformdiag definition --- NDTensors/src/diag/diagtensor.jl | 4 ---- 1 file changed, 4 deletions(-) diff --git a/NDTensors/src/diag/diagtensor.jl b/NDTensors/src/diag/diagtensor.jl index b2534b53d1..90035fac0a 100644 --- a/NDTensors/src/diag/diagtensor.jl +++ b/NDTensors/src/diag/diagtensor.jl @@ -35,10 +35,6 @@ function DiagonalArrays.diagview(T::NonuniformDiagTensor) return data(T) end -function DiagonalArrays.diagview(T::UniformDiagTensor) - return fill(getdiagindex(T, 1), diaglength(T)) -end - ## Should this go in dense.jl or here since its related to diag? function DiagonalArrays.diagview(T::DenseTensor) return diagview(array(T)) From 7d1ef42779fafe5ecf643745f7c6767bc9aa4e85 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 13 Jun 2024 19:55:22 -0400 Subject: [PATCH 17/41] remove comment --- NDTensors/src/diag/diagtensor.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/NDTensors/src/diag/diagtensor.jl b/NDTensors/src/diag/diagtensor.jl index 90035fac0a..8c32f408dc 100644 --- a/NDTensors/src/diag/diagtensor.jl +++ b/NDTensors/src/diag/diagtensor.jl @@ -35,7 +35,6 @@ function DiagonalArrays.diagview(T::NonuniformDiagTensor) return data(T) end -## Should this go in dense.jl or here since its related to diag? function DiagonalArrays.diagview(T::DenseTensor) return diagview(array(T)) end From ad4d96ca985c3e6a7f9dad2c324de8fa1320a1a4 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 13 Jun 2024 19:56:51 -0400 Subject: [PATCH 18/41] Move dense diagview to densetensor --- NDTensors/src/dense/densetensor.jl | 6 ++++++ NDTensors/src/diag/diagtensor.jl | 4 ---- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/NDTensors/src/dense/densetensor.jl b/NDTensors/src/dense/densetensor.jl index d73234a6a9..8e938cc387 100644 --- a/NDTensors/src/dense/densetensor.jl +++ b/NDTensors/src/dense/densetensor.jl @@ -68,6 +68,12 @@ convert(::Type{Array}, T::DenseTensor) = reshape(data(storage(T)), dims(inds(T)) # Useful for using Base Array functions array(T::DenseTensor) = convert(Array, T) +using .DiagonalArrays: DiagonalArrays, diagview + +function DiagonalArrays.diagview(T::DenseTensor) + return diagview(array(T)) +end + function Array{ElT,N}(T::DenseTensor{ElT,N}) where {ElT,N} return copy(array(T)) end diff --git a/NDTensors/src/diag/diagtensor.jl b/NDTensors/src/diag/diagtensor.jl index 8c32f408dc..34657ee94e 100644 --- a/NDTensors/src/diag/diagtensor.jl +++ b/NDTensors/src/diag/diagtensor.jl @@ -35,10 +35,6 @@ function DiagonalArrays.diagview(T::NonuniformDiagTensor) return data(T) end -function DiagonalArrays.diagview(T::DenseTensor) - return diagview(array(T)) -end - function zeros(tensortype::Type{<:DiagTensor}, inds) return tensor(generic_zeros(storagetype(tensortype), mindim(inds)), inds) end From 06d950314765ae909b7708bcab704954adaf5474 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 13 Jun 2024 19:57:31 -0400 Subject: [PATCH 19/41] typo --- NDTensors/src/diag/diagtensor.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NDTensors/src/diag/diagtensor.jl b/NDTensors/src/diag/diagtensor.jl index 34657ee94e..b6cea187fd 100644 --- a/NDTensors/src/diag/diagtensor.jl +++ b/NDTensors/src/diag/diagtensor.jl @@ -193,7 +193,7 @@ end function permutedims!( R::DenseTensor{ElR,N}, T::DiagTensor{ElT,N}, perm::NTuple{N,Int}, f::Function=(r, t) -> t ) where {ElR,ElT,N} - diagview(array(R)) .= f.(diagview(array(R)), diagview(T)) + diagview(R) .= f.(diagview(R), diagview(T)) return R end From 2140eb4457d2a9a66583674271106c2a57ce8238 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 13 Jun 2024 19:58:58 -0400 Subject: [PATCH 20/41] use diagview over data --- NDTensors/src/tensor/tensor.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NDTensors/src/tensor/tensor.jl b/NDTensors/src/tensor/tensor.jl index 16e90ad14b..cfc12bd728 100644 --- a/NDTensors/src/tensor/tensor.jl +++ b/NDTensors/src/tensor/tensor.jl @@ -369,7 +369,7 @@ function diag(tensor::Tensor) tensordiag = NDTensors.similar( dense(typeof(tensor)), eltype(tensor), (diaglength(tensor),) ) - data(tensordiag) .= diagview(tensor) + diagview(tensordiag) .= diagview(tensor) return tensordiag end From 51323fced1b3b3cc15195bf35f82e45dfadeb82f Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 13 Jun 2024 19:59:53 -0400 Subject: [PATCH 21/41] remove diaview --- NDTensors/src/tensor/tensor.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NDTensors/src/tensor/tensor.jl b/NDTensors/src/tensor/tensor.jl index cfc12bd728..a255aa82ef 100644 --- a/NDTensors/src/tensor/tensor.jl +++ b/NDTensors/src/tensor/tensor.jl @@ -369,7 +369,7 @@ function diag(tensor::Tensor) tensordiag = NDTensors.similar( dense(typeof(tensor)), eltype(tensor), (diaglength(tensor),) ) - diagview(tensordiag) .= diagview(tensor) + tensordiag .= diagview(tensor) return tensordiag end From d8cb8bdcded85a5e6c93c71f9b133b6eda3d7d55 Mon Sep 17 00:00:00 2001 From: kmp5VT Date: Fri, 14 Jun 2024 13:47:06 -0400 Subject: [PATCH 22/41] array necessary here because .= fails for CUDA gpu --- NDTensors/src/tensor/tensor.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NDTensors/src/tensor/tensor.jl b/NDTensors/src/tensor/tensor.jl index a255aa82ef..4ff25c46dd 100644 --- a/NDTensors/src/tensor/tensor.jl +++ b/NDTensors/src/tensor/tensor.jl @@ -369,7 +369,7 @@ function diag(tensor::Tensor) tensordiag = NDTensors.similar( dense(typeof(tensor)), eltype(tensor), (diaglength(tensor),) ) - tensordiag .= diagview(tensor) + array(tensordiag) .= diagview(tensor) return tensordiag end From 4f5e96c3fea22528133ba504b0ebce4130f81335 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Mon, 17 Jun 2024 15:48:05 -0400 Subject: [PATCH 23/41] Working on diag function --- NDTensors/src/blocksparse/blocksparsetensor.jl | 12 +++++++++++- NDTensors/src/tensor/tensor.jl | 9 ++++++--- NDTensors/test/test_blocksparse.jl | 3 +++ 3 files changed, 20 insertions(+), 4 deletions(-) diff --git a/NDTensors/src/blocksparse/blocksparsetensor.jl b/NDTensors/src/blocksparse/blocksparsetensor.jl index 6ff06ab0ca..c0f688ee4a 100644 --- a/NDTensors/src/blocksparse/blocksparsetensor.jl +++ b/NDTensors/src/blocksparse/blocksparsetensor.jl @@ -256,7 +256,7 @@ end # Returns the offset of the new block added. # XXX rename to insertblock!, no need to return offset using .TypeParameterAccessors: unwrap_array_type -using .Expose: expose +using .Expose: Exposed, expose, unexpose function insertblock_offset!(T::BlockSparseTensor{ElT,N}, newblock::Block{N}) where {ElT,N} newdim = blockdim(T, newblock) newoffset = nnz(T) @@ -356,6 +356,16 @@ function dense(T::TensorT) where {TensorT<:BlockSparseTensor} return tensor(Dense(r), inds(T)) end +function diag(ETensor::Exposed{<:AbstractArray, BlockSparseTensor}) + tensor = unexpose(ETensor) + tensordiag = NDTensors.similar( + dense(typeof(tensor)), eltype(tensor), (diaglength(tensor),) + ) + for j in 1:diaglength(tensor) + @inbounds tensor_diag[j] = getdiagindex(tensor, j) + end + return tensordiag +end # # Operations # diff --git a/NDTensors/src/tensor/tensor.jl b/NDTensors/src/tensor/tensor.jl index 4ff25c46dd..55e683427a 100644 --- a/NDTensors/src/tensor/tensor.jl +++ b/NDTensors/src/tensor/tensor.jl @@ -361,10 +361,13 @@ function getdiagindex(T::Tensor{<:Number,N}, ind::Int) where {N} return getindex(T, CartesianIndex(ntuple(_ -> ind, Val(N)))) end -using .DiagonalArrays: diagview +using .Expose: Exposed, expose, unexpose # TODO: add support for off-diagonals, return # block sparse vector instead of dense. -function diag(tensor::Tensor) +diag(tensor::Tensor) = diag(expose(tensor)) + +function diag(Etensor::Exposed) + tensor = unexpose(Etensor) ## d = NDTensors.similar(T, ElT, (diaglength(T),)) tensordiag = NDTensors.similar( dense(typeof(tensor)), eltype(tensor), (diaglength(tensor),) @@ -372,7 +375,7 @@ function diag(tensor::Tensor) array(tensordiag) .= diagview(tensor) return tensordiag end - + """ setdiagindex! diff --git a/NDTensors/test/test_blocksparse.jl b/NDTensors/test/test_blocksparse.jl index 3ed4379eba..427821286d 100644 --- a/NDTensors/test/test_blocksparse.jl +++ b/NDTensors/test/test_blocksparse.jl @@ -10,6 +10,7 @@ using NDTensors: blockview, data, dense, + diag, dims, eachnzblock, inds, @@ -52,6 +53,8 @@ using Test: @test, @test_throws, @testset @test isblocknz(A, (1, 2)) @test !isblocknz(A, (1, 1)) @test !isblocknz(A, (2, 2)) + dA = diag(A) + @allowscalar dA ≈ diag(dense(A)) # Test different ways of getting nnz @test nnz(blockoffsets(A), inds(A)) == nnz(A) From b0f15357087dd609eb6a673669aef7880b68d46c Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Mon, 17 Jun 2024 16:53:59 -0400 Subject: [PATCH 24/41] Attempt to update diag using expose --- .../NDTensorsGPUArraysCoreExt.jl | 1 + NDTensors/ext/NDTensorsGPUArraysCoreExt/diag.jl | 11 +++++++++++ NDTensors/src/blocksparse/blocksparsetensor.jl | 4 ++-- NDTensors/src/tensor/tensor.jl | 6 +++--- 4 files changed, 17 insertions(+), 5 deletions(-) create mode 100644 NDTensors/ext/NDTensorsGPUArraysCoreExt/diag.jl diff --git a/NDTensors/ext/NDTensorsGPUArraysCoreExt/NDTensorsGPUArraysCoreExt.jl b/NDTensors/ext/NDTensorsGPUArraysCoreExt/NDTensorsGPUArraysCoreExt.jl index 1ead107b45..7bd9f2a256 100644 --- a/NDTensors/ext/NDTensorsGPUArraysCoreExt/NDTensorsGPUArraysCoreExt.jl +++ b/NDTensors/ext/NDTensorsGPUArraysCoreExt/NDTensorsGPUArraysCoreExt.jl @@ -1,3 +1,4 @@ module NDTensorsGPUArraysCoreExt include("contract.jl") +include("diag.jl") end diff --git a/NDTensors/ext/NDTensorsGPUArraysCoreExt/diag.jl b/NDTensors/ext/NDTensorsGPUArraysCoreExt/diag.jl new file mode 100644 index 0000000000..5a82bf2caf --- /dev/null +++ b/NDTensors/ext/NDTensorsGPUArraysCoreExt/diag.jl @@ -0,0 +1,11 @@ +using GPUArraysCore: AbstractGPUArray +using NDTensors: NDTensors, BlockSparseTensor, dense, diag +using NDTensors.Expose: Exposed, unexpose + +## TODO to circumvent issues with blocksparse and scalar indexing +## convert blocksparse GPU tensors to dense tensors and call diag +## copying will probably have some impact on timing but this code +## currently isn't used in the main code, just in tests. +function NDTensors.diag(ETensor::Exposed{<:AbstractGPUArray,<:BlockSparseTensor}) + return diag(dense(unexpose(ETensor))) +end diff --git a/NDTensors/src/blocksparse/blocksparsetensor.jl b/NDTensors/src/blocksparse/blocksparsetensor.jl index c0f688ee4a..8a8200975b 100644 --- a/NDTensors/src/blocksparse/blocksparsetensor.jl +++ b/NDTensors/src/blocksparse/blocksparsetensor.jl @@ -356,13 +356,13 @@ function dense(T::TensorT) where {TensorT<:BlockSparseTensor} return tensor(Dense(r), inds(T)) end -function diag(ETensor::Exposed{<:AbstractArray, BlockSparseTensor}) +function diag(ETensor::Exposed{<:AbstractArray,<:BlockSparseTensor}) tensor = unexpose(ETensor) tensordiag = NDTensors.similar( dense(typeof(tensor)), eltype(tensor), (diaglength(tensor),) ) for j in 1:diaglength(tensor) - @inbounds tensor_diag[j] = getdiagindex(tensor, j) + @inbounds tensordiag[j] = getdiagindex(tensor, j) end return tensordiag end diff --git a/NDTensors/src/tensor/tensor.jl b/NDTensors/src/tensor/tensor.jl index 55e683427a..f35878210c 100644 --- a/NDTensors/src/tensor/tensor.jl +++ b/NDTensors/src/tensor/tensor.jl @@ -366,8 +366,8 @@ using .Expose: Exposed, expose, unexpose # block sparse vector instead of dense. diag(tensor::Tensor) = diag(expose(tensor)) -function diag(Etensor::Exposed) - tensor = unexpose(Etensor) +function diag(ETensor::Exposed) + tensor = unexpose(ETensor) ## d = NDTensors.similar(T, ElT, (diaglength(T),)) tensordiag = NDTensors.similar( dense(typeof(tensor)), eltype(tensor), (diaglength(tensor),) @@ -375,7 +375,7 @@ function diag(Etensor::Exposed) array(tensordiag) .= diagview(tensor) return tensordiag end - + """ setdiagindex! From 1fac29164ba7a380ff23f2d8e87933f90266a24f Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Mon, 17 Jun 2024 19:48:01 -0400 Subject: [PATCH 25/41] Update if statement --- NDTensors/test/test_diag.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NDTensors/test/test_diag.jl b/NDTensors/test/test_diag.jl index d4f6cbf188..0e6764e8fc 100644 --- a/NDTensors/test/test_diag.jl +++ b/NDTensors/test/test_diag.jl @@ -48,7 +48,7 @@ using LinearAlgebra: dot # This if statement corresponds to the reported bug: # https://github.com/JuliaGPU/Metal.jl/issues/364 - if (dev == NDTensors.mtl && elt != ComplexF32) + if !(dev == NDTensors.mtl && elt === ComplexF32) S = permutedims(dev(D), (1, 2), sqrt) @allowscalar begin for i in 1:diaglength(S) From 7e2ce770ff9308f353265dbddb28b8bd2d1c4192 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Wed, 19 Jun 2024 10:25:42 -0400 Subject: [PATCH 26/41] Move map_diag! to NDTensors --- NDTensors/src/tensor/tensor.jl | 7 +++++++ src/tensor_operations/matrix_algebra.jl | 12 +++--------- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/NDTensors/src/tensor/tensor.jl b/NDTensors/src/tensor/tensor.jl index f35878210c..e9d2088c5f 100644 --- a/NDTensors/src/tensor/tensor.jl +++ b/NDTensors/src/tensor/tensor.jl @@ -386,6 +386,13 @@ function setdiagindex!(T::Tensor{<:Number,N}, val, ind::Int) where {N} return T end +## TODO This fails for blocksparse because there is no diagview in blocksparse +function map_diag!(f::Function, t_destination::Tensor, t_source::Tensor) + diagview(t_destination) .= f.(diagview(t_source)) + return t_destination +end +map_diag(f::Function, t::Tensor) = map_diag!(f, copy(t), t) + # # Some generic contraction functionality # diff --git a/src/tensor_operations/matrix_algebra.jl b/src/tensor_operations/matrix_algebra.jl index 56660c8c64..61c7ca49a3 100644 --- a/src/tensor_operations/matrix_algebra.jl +++ b/src/tensor_operations/matrix_algebra.jl @@ -1,3 +1,5 @@ +using NDTensors: map_diag! + # Fix for AD function _tr(T::ITensor; plev::Pair{Int,Int}=0 => 1, tags::Pair=ts"" => ts"") trpairs = indpairs(T; plev=plev, tags=tags) @@ -86,15 +88,7 @@ function exp(A::ITensor; kwargs...) return exp(A, Lis, Ris; kwargs...) end -function map_diag!(f::Function, it_destination::ITensor, it_source::ITensor) +function NDTensors.map_diag!(f::Function, it_destination::ITensor, it_source::ITensor) return itensor(map_diag!(f, tensor(it_destination), tensor(it_source))) end map_diag(f::Function, it::ITensor) = map_diag!(f, copy(it), it) - -function map_diag!(f::Function, t_destination::Tensor, t_source::Tensor) - for i in 1:diaglength(t_destination) - NDTensors.setdiagindex!(t_destination, f(NDTensors.getdiagindex(t_source, i)), i) - end - return t_destination -end -map_diag(f::Function, t::Tensor) = map_diag!(f, copy(t), t) From 6ed349a3718f6e09df7707a0f26c46d0b71abeb5 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Wed, 19 Jun 2024 11:32:26 -0400 Subject: [PATCH 27/41] Don't use NDTensors map_diag yet --- src/tensor_operations/matrix_algebra.jl | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/src/tensor_operations/matrix_algebra.jl b/src/tensor_operations/matrix_algebra.jl index 61c7ca49a3..fdc2401756 100644 --- a/src/tensor_operations/matrix_algebra.jl +++ b/src/tensor_operations/matrix_algebra.jl @@ -1,5 +1,3 @@ -using NDTensors: map_diag! - # Fix for AD function _tr(T::ITensor; plev::Pair{Int,Int}=0 => 1, tags::Pair=ts"" => ts"") trpairs = indpairs(T; plev=plev, tags=tags) @@ -88,7 +86,15 @@ function exp(A::ITensor; kwargs...) return exp(A, Lis, Ris; kwargs...) end -function NDTensors.map_diag!(f::Function, it_destination::ITensor, it_source::ITensor) - return itensor(map_diag!(f, tensor(it_destination), tensor(it_source))) +function map_diag!(f::Function, it_destination::ITensor, it_source::ITensor) + return itensor(map_diag!(f, expose(tensor(it_destination)), expose(tensor(it_source)))) end map_diag(f::Function, it::ITensor) = map_diag!(f, copy(it), it) + +function map_diag!(f::Function, t_destination::Tensor, t_source::Tensor) + for i in 1:diaglength(t_destination) + NDTensors.setdiagindex!(t_destination, f(NDTensors.getdiagindex(t_source, i)), i) + end + return t_destination +end +map_diag(f::Function, t::Tensor) = map_diag!(f, copy(t), t) \ No newline at end of file From 5d049498d50040b05aa1e455701e2c17b8d8d6bc Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Wed, 19 Jun 2024 12:18:12 -0400 Subject: [PATCH 28/41] make map_diag expose based and create GPU version for blocksparse with allowscalar --- NDTensors/ext/NDTensorsGPUArraysCoreExt/diag.jl | 15 +++++++++++++-- NDTensors/src/blocksparse/blocksparsetensor.jl | 10 ++++++++++ NDTensors/src/tensor/tensor.jl | 7 +++---- 3 files changed, 26 insertions(+), 6 deletions(-) diff --git a/NDTensors/ext/NDTensorsGPUArraysCoreExt/diag.jl b/NDTensors/ext/NDTensorsGPUArraysCoreExt/diag.jl index 5a82bf2caf..d537c46421 100644 --- a/NDTensors/ext/NDTensorsGPUArraysCoreExt/diag.jl +++ b/NDTensors/ext/NDTensorsGPUArraysCoreExt/diag.jl @@ -1,5 +1,6 @@ -using GPUArraysCore: AbstractGPUArray -using NDTensors: NDTensors, BlockSparseTensor, dense, diag +using GPUArraysCore: @allowscalar, AbstractGPUArray +using NDTensors: NDTensors, BlockSparseTensor, dense, diag, map_diag! +using NDTensors.DiagonalArrays: diaglength using NDTensors.Expose: Exposed, unexpose ## TODO to circumvent issues with blocksparse and scalar indexing @@ -9,3 +10,13 @@ using NDTensors.Expose: Exposed, unexpose function NDTensors.diag(ETensor::Exposed{<:AbstractGPUArray,<:BlockSparseTensor}) return diag(dense(unexpose(ETensor))) end + +## TODO scalar indexing is slow here +function NDTensors.map_diag!(f::Function, exposed_t_destination::Exposed{<:AbstractGPUArray, <:BlockSparseTensor}, exposed_t_source::Exposed{<:AbstractGPUArray, <:BlockSparseTensor}) + t_destination = unexpose(exposed_t_destination) + t_source = unexpose(exposed_t_source) + @allowscalar for i in 1:diaglength(t_destination) + NDTensors.setdiagindex!(t_destination, f(NDTensors.getdiagindex(t_source, i)), i) + end + return t_destination +end \ No newline at end of file diff --git a/NDTensors/src/blocksparse/blocksparsetensor.jl b/NDTensors/src/blocksparse/blocksparsetensor.jl index 8a8200975b..99ef5c610f 100644 --- a/NDTensors/src/blocksparse/blocksparsetensor.jl +++ b/NDTensors/src/blocksparse/blocksparsetensor.jl @@ -366,6 +366,16 @@ function diag(ETensor::Exposed{<:AbstractArray,<:BlockSparseTensor}) end return tensordiag end + +## TODO currently this fails on GPU with scalar indexing +function map_diag!(f::Function, exposed_t_destination::Exposed{<:AbstractArray, <:BlockSparseTensor}, exposed_t_source::Exposed{<:AbstractArray, <:BlockSparseTensor}) + t_destination = unexpose(exposed_t_destination) + t_source = unexpose(exposed_t_source) + for i in 1:diaglength(t_destination) + NDTensors.setdiagindex!(t_destination, f(NDTensors.getdiagindex(t_source, i)), i) + end + return t_destination +end # # Operations # diff --git a/NDTensors/src/tensor/tensor.jl b/NDTensors/src/tensor/tensor.jl index e9d2088c5f..38e92063c6 100644 --- a/NDTensors/src/tensor/tensor.jl +++ b/NDTensors/src/tensor/tensor.jl @@ -386,12 +386,11 @@ function setdiagindex!(T::Tensor{<:Number,N}, val, ind::Int) where {N} return T end -## TODO This fails for blocksparse because there is no diagview in blocksparse -function map_diag!(f::Function, t_destination::Tensor, t_source::Tensor) - diagview(t_destination) .= f.(diagview(t_source)) +function map_diag!(f::Function, t_destination::Exposed, t_source::Exposed) + diagview(unexpose(t_destination)) .= f.(diagview(unexpose(t_source))) return t_destination end -map_diag(f::Function, t::Tensor) = map_diag!(f, copy(t), t) +map_diag(f::Function, t::Tensor) = map_diag!(f, expose(copy(t)), expose(t)) # # Some generic contraction functionality From 60f2abc1f6920d4a6014adba05ce47a33c859cff Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Wed, 19 Jun 2024 12:28:33 -0400 Subject: [PATCH 29/41] format --- NDTensors/ext/NDTensorsGPUArraysCoreExt/diag.jl | 8 ++++++-- NDTensors/src/blocksparse/blocksparsetensor.jl | 6 +++++- src/tensor_operations/matrix_algebra.jl | 2 +- 3 files changed, 12 insertions(+), 4 deletions(-) diff --git a/NDTensors/ext/NDTensorsGPUArraysCoreExt/diag.jl b/NDTensors/ext/NDTensorsGPUArraysCoreExt/diag.jl index d537c46421..1073bd2495 100644 --- a/NDTensors/ext/NDTensorsGPUArraysCoreExt/diag.jl +++ b/NDTensors/ext/NDTensorsGPUArraysCoreExt/diag.jl @@ -12,11 +12,15 @@ function NDTensors.diag(ETensor::Exposed{<:AbstractGPUArray,<:BlockSparseTensor} end ## TODO scalar indexing is slow here -function NDTensors.map_diag!(f::Function, exposed_t_destination::Exposed{<:AbstractGPUArray, <:BlockSparseTensor}, exposed_t_source::Exposed{<:AbstractGPUArray, <:BlockSparseTensor}) +function NDTensors.map_diag!( + f::Function, + exposed_t_destination::Exposed{<:AbstractGPUArray,<:BlockSparseTensor}, + exposed_t_source::Exposed{<:AbstractGPUArray,<:BlockSparseTensor}, +) t_destination = unexpose(exposed_t_destination) t_source = unexpose(exposed_t_source) @allowscalar for i in 1:diaglength(t_destination) NDTensors.setdiagindex!(t_destination, f(NDTensors.getdiagindex(t_source, i)), i) end return t_destination -end \ No newline at end of file +end diff --git a/NDTensors/src/blocksparse/blocksparsetensor.jl b/NDTensors/src/blocksparse/blocksparsetensor.jl index 99ef5c610f..77ed73b40b 100644 --- a/NDTensors/src/blocksparse/blocksparsetensor.jl +++ b/NDTensors/src/blocksparse/blocksparsetensor.jl @@ -368,7 +368,11 @@ function diag(ETensor::Exposed{<:AbstractArray,<:BlockSparseTensor}) end ## TODO currently this fails on GPU with scalar indexing -function map_diag!(f::Function, exposed_t_destination::Exposed{<:AbstractArray, <:BlockSparseTensor}, exposed_t_source::Exposed{<:AbstractArray, <:BlockSparseTensor}) +function map_diag!( + f::Function, + exposed_t_destination::Exposed{<:AbstractArray,<:BlockSparseTensor}, + exposed_t_source::Exposed{<:AbstractArray,<:BlockSparseTensor}, +) t_destination = unexpose(exposed_t_destination) t_source = unexpose(exposed_t_source) for i in 1:diaglength(t_destination) diff --git a/src/tensor_operations/matrix_algebra.jl b/src/tensor_operations/matrix_algebra.jl index fdc2401756..383a38acc4 100644 --- a/src/tensor_operations/matrix_algebra.jl +++ b/src/tensor_operations/matrix_algebra.jl @@ -97,4 +97,4 @@ function map_diag!(f::Function, t_destination::Tensor, t_source::Tensor) end return t_destination end -map_diag(f::Function, t::Tensor) = map_diag!(f, copy(t), t) \ No newline at end of file +map_diag(f::Function, t::Tensor) = map_diag!(f, copy(t), t) From 94337986405659bf3f341a5db06d88930e022853 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Wed, 19 Jun 2024 14:00:57 -0400 Subject: [PATCH 30/41] remove expose --- src/tensor_operations/matrix_algebra.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/tensor_operations/matrix_algebra.jl b/src/tensor_operations/matrix_algebra.jl index 383a38acc4..56660c8c64 100644 --- a/src/tensor_operations/matrix_algebra.jl +++ b/src/tensor_operations/matrix_algebra.jl @@ -87,7 +87,7 @@ function exp(A::ITensor; kwargs...) end function map_diag!(f::Function, it_destination::ITensor, it_source::ITensor) - return itensor(map_diag!(f, expose(tensor(it_destination)), expose(tensor(it_source)))) + return itensor(map_diag!(f, tensor(it_destination), tensor(it_source))) end map_diag(f::Function, it::ITensor) = map_diag!(f, copy(it), it) From 3df285eb914a555c4d1d2f8595e9316c2b0ba735 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Wed, 19 Jun 2024 14:41:09 -0400 Subject: [PATCH 31/41] simplify dense definition for DiagTensor --- NDTensors/src/diag/diagtensor.jl | 50 +++++++++++++++++--------------- 1 file changed, 26 insertions(+), 24 deletions(-) diff --git a/NDTensors/src/diag/diagtensor.jl b/NDTensors/src/diag/diagtensor.jl index b6cea187fd..4eef6cea3c 100644 --- a/NDTensors/src/diag/diagtensor.jl +++ b/NDTensors/src/diag/diagtensor.jl @@ -112,31 +112,33 @@ end using .TypeParameterAccessors: unwrap_array_type # convert to Dense function dense(T::DiagTensor) - return dense(unwrap_array_type(T), T) -end - -# CPU version -function dense(::Type{<:Array}, T::DiagTensor) R = zeros(dense(typeof(T)), inds(T)) - for i in 1:diaglength(T) - setdiagindex!(R, getdiagindex(T, i), i) - end - return R -end - -# GPU version -function dense(::Type{<:AbstractArray}, T::DiagTensor) - D_cpu = dense(Array, cpu(T)) - return adapt(unwrap_array_type(T), D_cpu) -end - -# UniformDiag version -# TODO: Delete once new DiagonalArray is designed. -# TODO: This creates a tensor on CPU by default so may cause -# problems for GPU. -function dense(::Type{<:Number}, T::DiagTensor) - return dense(Tensor(Diag(fill(getdiagindex(T, 1), diaglength(T))), inds(T))) -end + diagview(R) .= diagview(T) + return R +end + +# # CPU version +# function dense(::Type{<:Array}, T::DiagTensor) +# R = zeros(dense(typeof(T)), inds(T)) +# for i in 1:diaglength(T) +# setdiagindex!(R, getdiagindex(T, i), i) +# end +# return R +# end + +# # GPU version +# function dense(::Type{<:AbstractArray}, T::DiagTensor) +# D_cpu = dense(Array, cpu(T)) +# return adapt(unwrap_array_type(T), D_cpu) +# end + +# # UniformDiag version +# # TODO: Delete once new DiagonalArray is designed. +# # TODO: This creates a tensor on CPU by default so may cause +# # problems for GPU. +# function dense(::Type{<:Number}, T::DiagTensor) +# return dense(Tensor(Diag(fill(getdiagindex(T, 1), diaglength(T))), inds(T))) +# end denseblocks(T::DiagTensor) = dense(T) From 8e492883b610a35e58ea198da078ce0075d8e929 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Wed, 19 Jun 2024 14:41:32 -0400 Subject: [PATCH 32/41] remove unused code --- NDTensors/src/diag/diagtensor.jl | 27 ++------------------------- 1 file changed, 2 insertions(+), 25 deletions(-) diff --git a/NDTensors/src/diag/diagtensor.jl b/NDTensors/src/diag/diagtensor.jl index 4eef6cea3c..bfde30a397 100644 --- a/NDTensors/src/diag/diagtensor.jl +++ b/NDTensors/src/diag/diagtensor.jl @@ -114,31 +114,8 @@ using .TypeParameterAccessors: unwrap_array_type function dense(T::DiagTensor) R = zeros(dense(typeof(T)), inds(T)) diagview(R) .= diagview(T) - return R -end - -# # CPU version -# function dense(::Type{<:Array}, T::DiagTensor) -# R = zeros(dense(typeof(T)), inds(T)) -# for i in 1:diaglength(T) -# setdiagindex!(R, getdiagindex(T, i), i) -# end -# return R -# end - -# # GPU version -# function dense(::Type{<:AbstractArray}, T::DiagTensor) -# D_cpu = dense(Array, cpu(T)) -# return adapt(unwrap_array_type(T), D_cpu) -# end - -# # UniformDiag version -# # TODO: Delete once new DiagonalArray is designed. -# # TODO: This creates a tensor on CPU by default so may cause -# # problems for GPU. -# function dense(::Type{<:Number}, T::DiagTensor) -# return dense(Tensor(Diag(fill(getdiagindex(T, 1), diaglength(T))), inds(T))) -# end + return R +end denseblocks(T::DiagTensor) = dense(T) From caba6a4e81cc823ee6c95b6df9896b0b896de982 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 20 Jun 2024 09:24:26 -0400 Subject: [PATCH 33/41] Rename diag to blocksparsetensor.jl --- .../ext/NDTensorsGPUArraysCoreExt/NDTensorsGPUArraysCoreExt.jl | 2 +- .../NDTensorsGPUArraysCoreExt/{diag.jl => blocksparsetensor.jl} | 0 2 files changed, 1 insertion(+), 1 deletion(-) rename NDTensors/ext/NDTensorsGPUArraysCoreExt/{diag.jl => blocksparsetensor.jl} (100%) diff --git a/NDTensors/ext/NDTensorsGPUArraysCoreExt/NDTensorsGPUArraysCoreExt.jl b/NDTensors/ext/NDTensorsGPUArraysCoreExt/NDTensorsGPUArraysCoreExt.jl index 7bd9f2a256..c9e183cd52 100644 --- a/NDTensors/ext/NDTensorsGPUArraysCoreExt/NDTensorsGPUArraysCoreExt.jl +++ b/NDTensors/ext/NDTensorsGPUArraysCoreExt/NDTensorsGPUArraysCoreExt.jl @@ -1,4 +1,4 @@ module NDTensorsGPUArraysCoreExt include("contract.jl") -include("diag.jl") +include("blocksparsetensor.jl") end diff --git a/NDTensors/ext/NDTensorsGPUArraysCoreExt/diag.jl b/NDTensors/ext/NDTensorsGPUArraysCoreExt/blocksparsetensor.jl similarity index 100% rename from NDTensors/ext/NDTensorsGPUArraysCoreExt/diag.jl rename to NDTensors/ext/NDTensorsGPUArraysCoreExt/blocksparsetensor.jl From ea52a82b185fcc61691854708d38df71515a536c Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 20 Jun 2024 09:26:26 -0400 Subject: [PATCH 34/41] Try forcing Pkg to update the registry to fix ci issue --- jenkins/Jenkinsfile | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/jenkins/Jenkinsfile b/jenkins/Jenkinsfile index f1254db10f..c50a7c874c 100644 --- a/jenkins/Jenkinsfile +++ b/jenkins/Jenkinsfile @@ -27,7 +27,7 @@ pipeline { } steps { sh ''' - julia -e 'using Pkg; Pkg.activate(temp=true); Pkg.develop(path="./NDTensors"); Pkg.develop(path="."); Pkg.test("NDTensors"; test_args=["cuda"])' + julia -e 'using Pkg; Pkg.Registry.update(); Pkg.update(); Pkg.activate(temp=true); Pkg.develop(path="./NDTensors"); Pkg.develop(path="."); Pkg.test("NDTensors"; test_args=["cuda"])' ''' } } @@ -51,7 +51,7 @@ pipeline { } steps { sh ''' - julia -e 'using Pkg; Pkg.activate(temp=true); Pkg.develop(path="./NDTensors"); Pkg.develop(path="."); Pkg.test("NDTensors"; test_args=["cuda"])' + julia -e 'using Pkg; Pkg.Registry.update(); Pkg.update(); Pkg.activate(temp=true); Pkg.develop(path="./NDTensors"); Pkg.develop(path="."); Pkg.test("NDTensors"; test_args=["cuda"])' ''' } } @@ -75,7 +75,7 @@ pipeline { } steps { sh ''' - julia -e 'using Pkg; Pkg.activate(temp=true); Pkg.develop(path="./NDTensors"); Pkg.develop(path="."); Pkg.test("NDTensors"; test_args=["cutensor"])' + julia -e 'using Pkg; Pkg.Registry.update(); Pkg.update(); Pkg.activate(temp=true); Pkg.develop(path="./NDTensors"); Pkg.develop(path="."); Pkg.test("NDTensors"; test_args=["cutensor"])' ''' } } From a48ddb0543b54b4827710b815e2e65dc74b04ca6 Mon Sep 17 00:00:00 2001 From: kmp5VT Date: Thu, 20 Jun 2024 10:22:56 -0400 Subject: [PATCH 35/41] Use approx because of numerical noise --- NDTensors/test/test_diag.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NDTensors/test/test_diag.jl b/NDTensors/test/test_diag.jl index 0e6764e8fc..be4596effe 100644 --- a/NDTensors/test/test_diag.jl +++ b/NDTensors/test/test_diag.jl @@ -52,7 +52,7 @@ using LinearAlgebra: dot S = permutedims(dev(D), (1, 2), sqrt) @allowscalar begin for i in 1:diaglength(S) - @test S[i, i] == sqrt(D[i, i]) + @test S[i, i] ≈ sqrt(D[i, i]) end end end From f5c2d399767205f324b2fc467ccf2411891f6791 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 20 Jun 2024 18:35:46 -0400 Subject: [PATCH 36/41] Add map_diag tests --- NDTensors/src/tensor/tensor.jl | 6 +++--- NDTensors/test/test_blocksparse.jl | 7 ++++++- NDTensors/test/test_dense.jl | 5 +++++ NDTensors/test/test_diag.jl | 11 +++++++++++ 4 files changed, 25 insertions(+), 4 deletions(-) diff --git a/NDTensors/src/tensor/tensor.jl b/NDTensors/src/tensor/tensor.jl index 38e92063c6..ac6fee9a6d 100644 --- a/NDTensors/src/tensor/tensor.jl +++ b/NDTensors/src/tensor/tensor.jl @@ -386,9 +386,9 @@ function setdiagindex!(T::Tensor{<:Number,N}, val, ind::Int) where {N} return T end -function map_diag!(f::Function, t_destination::Exposed, t_source::Exposed) - diagview(unexpose(t_destination)) .= f.(diagview(unexpose(t_source))) - return t_destination +function map_diag!(f::Function, exposed_t_destination::Exposed, exposed_t_source::Exposed) + diagview(unexpose(exposed_t_destination)) .= f.(diagview(unexpose(exposed_t_source))) + return unexpose(exposed_t_destination) end map_diag(f::Function, t::Tensor) = map_diag!(f, expose(copy(t)), expose(t)) diff --git a/NDTensors/test/test_blocksparse.jl b/NDTensors/test/test_blocksparse.jl index 427821286d..8a0753c4ba 100644 --- a/NDTensors/test/test_blocksparse.jl +++ b/NDTensors/test/test_blocksparse.jl @@ -11,6 +11,7 @@ using NDTensors: data, dense, diag, + diaglength, dims, eachnzblock, inds, @@ -54,7 +55,7 @@ using Test: @test, @test_throws, @testset @test !isblocknz(A, (1, 1)) @test !isblocknz(A, (2, 2)) dA = diag(A) - @allowscalar dA ≈ diag(dense(A)) + @test @allowscalar dA ≈ diag(dense(A)) # Test different ways of getting nnz @test nnz(blockoffsets(A), inds(A)) == nnz(A) @@ -107,6 +108,10 @@ using Test: @test, @test_throws, @testset @allowscalar for I in eachindex(C) @test C[I] == A[I] + B[I] end + Cp = NDTensors.map_diag((i -> i* 2), C) + @allowscalar for i in 1:diaglength(Cp) + @test Cp[i,i] == 2.0 * C[i,i] + end Ap = permutedims(A, (2, 1)) diff --git a/NDTensors/test/test_dense.jl b/NDTensors/test/test_dense.jl index 78368d2841..c5c1c8a231 100644 --- a/NDTensors/test/test_dense.jl +++ b/NDTensors/test/test_dense.jl @@ -48,6 +48,10 @@ NDTensors.dim(i::MyInd) = i.dim randn!(B) C = copy(A) C = permutedims!!(C, B, (1, 2), +) + Cp = NDTensors.map_diag((i->i * 2), C) + @allowscalar for i in 1:diaglength(Cp) + @test Cp[i,i] == 2.0 * C[i,i] + end Ap = permutedims(A, (2, 1)) @allowscalar begin @@ -307,3 +311,4 @@ end nothing end +w \ No newline at end of file diff --git a/NDTensors/test/test_diag.jl b/NDTensors/test/test_diag.jl index be4596effe..1d5e477ffb 100644 --- a/NDTensors/test/test_diag.jl +++ b/NDTensors/test/test_diag.jl @@ -34,6 +34,12 @@ using LinearAlgebra: dot D = Tensor(Diag(1), (2, 2)) @test norm(D) == √2 d = 3 + ## TODO this fails because uniform diag tensors are immutable + #S = NDTensors.map_diag((i->i * 2), dev(D)) + # @allowscalar for i in 1:diaglength(S) + # @test S[i,i] == 2.0 * D[i,i] + # end + vr = rand(elt, d) D = dev(tensor(Diag(vr), (d, d))) Da = Array(D) @@ -56,6 +62,11 @@ using LinearAlgebra: dot end end end + S = NDTensors.map_diag((i->i * 2), dev(D)) + @allowscalar for i in 1:diaglength(S) + @test S[i,i] == 2.0 * D[i,i] + end + # Regression test for https://github.com/ITensor/ITensors.jl/issues/1199 S = dev(tensor(Diag(randn(elt, 2)), (2, 2))) From d0dcb4dd89a5bf3f83b0dc9c11065d5b287feaaa Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 20 Jun 2024 18:36:38 -0400 Subject: [PATCH 37/41] format --- NDTensors/test/test_blocksparse.jl | 4 ++-- NDTensors/test/test_dense.jl | 6 +++--- NDTensors/test/test_diag.jl | 7 +++---- 3 files changed, 8 insertions(+), 9 deletions(-) diff --git a/NDTensors/test/test_blocksparse.jl b/NDTensors/test/test_blocksparse.jl index 8a0753c4ba..67a0d17144 100644 --- a/NDTensors/test/test_blocksparse.jl +++ b/NDTensors/test/test_blocksparse.jl @@ -108,9 +108,9 @@ using Test: @test, @test_throws, @testset @allowscalar for I in eachindex(C) @test C[I] == A[I] + B[I] end - Cp = NDTensors.map_diag((i -> i* 2), C) + Cp = NDTensors.map_diag((i -> i * 2), C) @allowscalar for i in 1:diaglength(Cp) - @test Cp[i,i] == 2.0 * C[i,i] + @test Cp[i, i] == 2.0 * C[i, i] end Ap = permutedims(A, (2, 1)) diff --git a/NDTensors/test/test_dense.jl b/NDTensors/test/test_dense.jl index c5c1c8a231..6ecd285e67 100644 --- a/NDTensors/test/test_dense.jl +++ b/NDTensors/test/test_dense.jl @@ -48,9 +48,9 @@ NDTensors.dim(i::MyInd) = i.dim randn!(B) C = copy(A) C = permutedims!!(C, B, (1, 2), +) - Cp = NDTensors.map_diag((i->i * 2), C) + Cp = NDTensors.map_diag((i -> i * 2), C) @allowscalar for i in 1:diaglength(Cp) - @test Cp[i,i] == 2.0 * C[i,i] + @test Cp[i, i] == 2.0 * C[i, i] end Ap = permutedims(A, (2, 1)) @@ -311,4 +311,4 @@ end nothing end -w \ No newline at end of file +w diff --git a/NDTensors/test/test_diag.jl b/NDTensors/test/test_diag.jl index 1d5e477ffb..c276de917b 100644 --- a/NDTensors/test/test_diag.jl +++ b/NDTensors/test/test_diag.jl @@ -39,7 +39,7 @@ using LinearAlgebra: dot # @allowscalar for i in 1:diaglength(S) # @test S[i,i] == 2.0 * D[i,i] # end - + vr = rand(elt, d) D = dev(tensor(Diag(vr), (d, d))) Da = Array(D) @@ -62,11 +62,10 @@ using LinearAlgebra: dot end end end - S = NDTensors.map_diag((i->i * 2), dev(D)) + S = NDTensors.map_diag((i -> i * 2), dev(D)) @allowscalar for i in 1:diaglength(S) - @test S[i,i] == 2.0 * D[i,i] + @test S[i, i] == 2.0 * D[i, i] end - # Regression test for https://github.com/ITensor/ITensors.jl/issues/1199 S = dev(tensor(Diag(randn(elt, 2)), (2, 2))) From ae19547f472d15ecc0700c6f831c229c6613cb65 Mon Sep 17 00:00:00 2001 From: Karl Pierce Date: Thu, 20 Jun 2024 19:07:50 -0400 Subject: [PATCH 38/41] remove w --- NDTensors/test/test_dense.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/NDTensors/test/test_dense.jl b/NDTensors/test/test_dense.jl index 6ecd285e67..4c7e274cd4 100644 --- a/NDTensors/test/test_dense.jl +++ b/NDTensors/test/test_dense.jl @@ -311,4 +311,3 @@ end nothing end -w From b2b0b195758407d7ba4eb64cf4ceb6704d2135de Mon Sep 17 00:00:00 2001 From: Matt Fishman Date: Thu, 20 Jun 2024 19:13:41 -0400 Subject: [PATCH 39/41] Format --- NDTensors/test/test_blocksparse.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/NDTensors/test/test_blocksparse.jl b/NDTensors/test/test_blocksparse.jl index 67a0d17144..b230effadc 100644 --- a/NDTensors/test/test_blocksparse.jl +++ b/NDTensors/test/test_blocksparse.jl @@ -108,9 +108,9 @@ using Test: @test, @test_throws, @testset @allowscalar for I in eachindex(C) @test C[I] == A[I] + B[I] end - Cp = NDTensors.map_diag((i -> i * 2), C) + Cp = NDTensors.map_diag(i -> 2 * i, C) @allowscalar for i in 1:diaglength(Cp) - @test Cp[i, i] == 2.0 * C[i, i] + @test Cp[i, i] == 2 * C[i, i] end Ap = permutedims(A, (2, 1)) From 255f66e93ec0f7bd1288ec19cce0b511cc20fa93 Mon Sep 17 00:00:00 2001 From: Matt Fishman Date: Thu, 20 Jun 2024 19:14:12 -0400 Subject: [PATCH 40/41] Format --- NDTensors/test/test_dense.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/NDTensors/test/test_dense.jl b/NDTensors/test/test_dense.jl index 4c7e274cd4..94c52f4132 100644 --- a/NDTensors/test/test_dense.jl +++ b/NDTensors/test/test_dense.jl @@ -48,9 +48,9 @@ NDTensors.dim(i::MyInd) = i.dim randn!(B) C = copy(A) C = permutedims!!(C, B, (1, 2), +) - Cp = NDTensors.map_diag((i -> i * 2), C) + Cp = NDTensors.map_diag(i -> 2 * i, C) @allowscalar for i in 1:diaglength(Cp) - @test Cp[i, i] == 2.0 * C[i, i] + @test Cp[i, i] == 2 * C[i, i] end Ap = permutedims(A, (2, 1)) From ab0728faaceea80df807f41486ad17795a3feee4 Mon Sep 17 00:00:00 2001 From: Matt Fishman Date: Thu, 20 Jun 2024 19:15:20 -0400 Subject: [PATCH 41/41] Format --- NDTensors/test/test_diag.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/NDTensors/test/test_diag.jl b/NDTensors/test/test_diag.jl index c276de917b..1876919117 100644 --- a/NDTensors/test/test_diag.jl +++ b/NDTensors/test/test_diag.jl @@ -62,9 +62,9 @@ using LinearAlgebra: dot end end end - S = NDTensors.map_diag((i -> i * 2), dev(D)) + S = NDTensors.map_diag(i -> 2 * i, dev(D)) @allowscalar for i in 1:diaglength(S) - @test S[i, i] == 2.0 * D[i, i] + @test S[i, i] == 2 * D[i, i] end # Regression test for https://github.com/ITensor/ITensors.jl/issues/1199