From c51a74484200c103aa974c28698a9735402b40e4 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Fri, 30 Jan 2026 06:39:48 -0500 Subject: [PATCH 1/6] Small fixes for upstream + CUDA --- ext/TensorKitCUDAExt/cutensormap.jl | 3 +++ src/tensors/tensor.jl | 11 ++++++++--- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/ext/TensorKitCUDAExt/cutensormap.jl b/ext/TensorKitCUDAExt/cutensormap.jl index 3274e654a..9f5c6f9ba 100644 --- a/ext/TensorKitCUDAExt/cutensormap.jl +++ b/ext/TensorKitCUDAExt/cutensormap.jl @@ -138,6 +138,9 @@ function Base.promote_rule( return CuTensorMap{T, S, N₁, N₂} end +TensorKit.promote_storagetype(TA::CuArray{T, N, M}, TB::CuArray{T, N}) where {T, N, M} = CuArray{T, N, M} +TensorKit.promote_storagetype(TA::CuArray{T, N, M}, TB::CuArray{T, N, M}) where {T, N, M} = CuArray{T, N, M} + # CuTensorMap exponentation: function TensorKit.exp!(t::CuTensorMap) domain(t) == codomain(t) || diff --git a/src/tensors/tensor.jl b/src/tensors/tensor.jl index 615c67775..7680ee121 100644 --- a/src/tensors/tensor.jl +++ b/src/tensors/tensor.jl @@ -306,9 +306,14 @@ for (fname, felt) in ((:zeros, :zero), (:ones, :one)) return Base.$fname(codomain ← domain) end function Base.$fname( - ::Type{TorA}, codomain::TensorSpace{S}, domain::TensorSpace{S} = one(codomain) - ) where {TorA, S <: IndexSpace} - return Base.$fname(TorA, codomain ← domain) + ::Type{T}, codomain::TensorSpace{S}, domain::TensorSpace{S} = one(codomain) + ) where {T <: Number, S <: IndexSpace} + return Base.$fname(T, codomain ← domain) + end + function Base.$fname( + ::Type{TA}, codomain::TensorSpace{S}, domain::TensorSpace{S} = one(codomain) + ) where {TA <: Array, S <: IndexSpace} + return Base.$fname(eltype(TA), codomain ← domain) end Base.$fname(V::TensorMapSpace) = Base.$fname(Float64, V) function Base.$fname(::Type{TorA}, V::TensorMapSpace) where {TorA} From f99e931c64b566f74bd7e57221a94f46b0436e1b Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Wed, 11 Feb 2026 14:26:21 -0500 Subject: [PATCH 2/6] Suggestion --- ext/TensorKitCUDAExt/cutensormap.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/ext/TensorKitCUDAExt/cutensormap.jl b/ext/TensorKitCUDAExt/cutensormap.jl index 9f5c6f9ba..158576ea6 100644 --- a/ext/TensorKitCUDAExt/cutensormap.jl +++ b/ext/TensorKitCUDAExt/cutensormap.jl @@ -138,8 +138,9 @@ function Base.promote_rule( return CuTensorMap{T, S, N₁, N₂} end -TensorKit.promote_storagetype(TA::CuArray{T, N, M}, TB::CuArray{T, N}) where {T, N, M} = CuArray{T, N, M} -TensorKit.promote_storagetype(TA::CuArray{T, N, M}, TB::CuArray{T, N, M}) where {T, N, M} = CuArray{T, N, M} +TensorKit.promote_storage_rule(::Type{CuArray{T, N, M}}, ::Type{<:CuArray{T, N}}) where {T, N, M} = + CuArray{T, N, M} + # CuTensorMap exponentation: function TensorKit.exp!(t::CuTensorMap) From 0f345f483bf420b69075cfa04e627089e0c4969a Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Wed, 11 Feb 2026 14:53:33 -0500 Subject: [PATCH 3/6] Use default memory --- ext/TensorKitCUDAExt/cutensormap.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ext/TensorKitCUDAExt/cutensormap.jl b/ext/TensorKitCUDAExt/cutensormap.jl index 158576ea6..f065c2ec1 100644 --- a/ext/TensorKitCUDAExt/cutensormap.jl +++ b/ext/TensorKitCUDAExt/cutensormap.jl @@ -138,8 +138,8 @@ function Base.promote_rule( return CuTensorMap{T, S, N₁, N₂} end -TensorKit.promote_storage_rule(::Type{CuArray{T, N, M}}, ::Type{<:CuArray{T, N}}) where {T, N, M} = - CuArray{T, N, M} +TensorKit.promote_storage_rule(::Type{CuArray{T, N}}, ::Type{<:CuArray{T, N}}) where {T, N} = + CuArray{T, N, CUDA.default_memory} # CuTensorMap exponentation: From 51fea051b74a2eb2beb9e011b7a955c7a4d19bc1 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Wed, 11 Feb 2026 15:26:24 -0500 Subject: [PATCH 4/6] Have Base.ones and zeros accept CuArray --- ext/TensorKitCUDAExt/cutensormap.jl | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/ext/TensorKitCUDAExt/cutensormap.jl b/ext/TensorKitCUDAExt/cutensormap.jl index f065c2ec1..d3287eef5 100644 --- a/ext/TensorKitCUDAExt/cutensormap.jl +++ b/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{TA}, codomain::TensorSpace{S}, domain::TensorSpace{S} = one(codomain) + ) where {TA <: CuArray, S <: IndexSpace} + return CUDA.$fname(eltype(TA), codomain ← domain) + end end end From 76c3a0c51f5c876cc22dc426dbf5d803eb9d2505 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 12 Feb 2026 01:22:38 -0500 Subject: [PATCH 5/6] Try to finagle ones and zeros again --- ext/TensorKitCUDAExt/cutensormap.jl | 5 ----- src/tensors/tensor.jl | 11 +++-------- test/cuda/tensors.jl | 8 ++++++++ test/tensors/tensors.jl | 9 +++++++++ 4 files changed, 20 insertions(+), 13 deletions(-) diff --git a/ext/TensorKitCUDAExt/cutensormap.jl b/ext/TensorKitCUDAExt/cutensormap.jl index d3287eef5..f065c2ec1 100644 --- a/ext/TensorKitCUDAExt/cutensormap.jl +++ b/ext/TensorKitCUDAExt/cutensormap.jl @@ -37,11 +37,6 @@ for (fname, felt) in ((:zeros, :zero), (:ones, :one)) fill!(t, $felt(T)) return t end - function Base.$fname( - ::Type{TA}, codomain::TensorSpace{S}, domain::TensorSpace{S} = one(codomain) - ) where {TA <: CuArray, S <: IndexSpace} - return CUDA.$fname(eltype(TA), codomain ← domain) - end end end diff --git a/src/tensors/tensor.jl b/src/tensors/tensor.jl index 7680ee121..615c67775 100644 --- a/src/tensors/tensor.jl +++ b/src/tensors/tensor.jl @@ -306,14 +306,9 @@ 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 <: Number, S <: IndexSpace} - return Base.$fname(T, codomain ← domain) - end - function Base.$fname( - ::Type{TA}, codomain::TensorSpace{S}, domain::TensorSpace{S} = one(codomain) - ) where {TA <: Array, S <: IndexSpace} - return Base.$fname(eltype(TA), codomain ← domain) + ::Type{TorA}, codomain::TensorSpace{S}, domain::TensorSpace{S} = one(codomain) + ) where {TorA, S <: IndexSpace} + return Base.$fname(TorA, codomain ← domain) end Base.$fname(V::TensorMapSpace) = Base.$fname(Float64, V) function Base.$fname(::Type{TorA}, V::TensorMapSpace) where {TorA} diff --git a/test/cuda/tensors.jl b/test/cuda/tensors.jl index a31015a4a..6df7afa6d 100644 --- a/test/cuda/tensors.jl +++ b/test/cuda/tensors.jl @@ -55,6 +55,14 @@ 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 (Base.ones, Base.zeros) + t = @constinferred f(CuVector{Float64}, W) + @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 f in (rand, randn) t = @constinferred f(CuVector{Float64, CUDA.DeviceMemory}, W) @test scalartype(t) == Float64 diff --git a/test/tensors/tensors.jl b/test/tensors/tensors.jl index 74cb9cfa0..813c25ea5 100644 --- a/test/tensors/tensors.jl +++ b/test/tensors/tensors.jl @@ -44,6 +44,15 @@ for V in spacelist @test space(t) == (W ← one(W)) @test domain(t) == one(W) @test typeof(t) == TensorMap{T, spacetype(t), 5, 0, Vector{T}} + # Array type input + t = @constinferred zeros(Vector{T}, W) + @test @constinferred(hash(t)) == hash(deepcopy(t)) + @test scalartype(t) == T + @test norm(t) == 0 + @test codomain(t) == W + @test space(t) == (W ← one(W)) + @test domain(t) == one(W) + @test typeof(t) == TensorMap{T, spacetype(t), 5, 0, Vector{T}} # blocks bs = @constinferred blocks(t) if !isempty(blocksectors(t)) # multifusion space ending on module gives empty data From b6189886ae413115866097cbb4610f283d342b6f Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 12 Feb 2026 03:10:52 -0500 Subject: [PATCH 6/6] Fix CUDA test --- test/cuda/tensors.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/cuda/tensors.jl b/test/cuda/tensors.jl index 6df7afa6d..0fad13473 100644 --- a/test/cuda/tensors.jl +++ b/test/cuda/tensors.jl @@ -56,7 +56,7 @@ for V in spacelist @test typeof(t) == TensorMap{Float64, spacetype(t), 5, 0, CuVector{Float64, CUDA.DeviceMemory}} end for f in (Base.ones, Base.zeros) - t = @constinferred f(CuVector{Float64}, W) + t = @constinferred f(CuVector{Float64, CUDA.DeviceMemory}, W) @test scalartype(t) == Float64 @test codomain(t) == W @test space(t) == (W ← one(W))