Try to finagle ones and zeros again

Author: kshyatt

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
 

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}

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

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