ones function is not type-stable for CUDA

[bgctw]

On a machine with CUDA GPU, the following minimal example returns error
return type CUDA.CuArray{Float32, 1, CUDA.DeviceMemory} does not match inferred return type CUDA.CuArray{Float32, 1}

#Pkg.activate(;temp=true)
#Pkg.add(["CUDA","cuDNN","MLDataDevices","Test","KernelAbstractions"])
import CUDA, cuDNN
using MLDataDevices
using Test
using KernelAbstractions

gdev = gpu_device()
x = rand(2,3)
xd = gdev(x)
backend = @inferred get_backend(xd)
onevec = @inferred KernelAbstractions.ones(backend, eltype(xd), 2)

With Julia 1.11.4 and installed packages

  [052768ef] CUDA v5.8.3
  [63c18a36] KernelAbstractions v0.9.38
  [7e8f7934] MLDataDevices v1.11.2
  [02a925ec] cuDNN v1.4.5
  [8dfed614] Test v1.11.0

[vchuravy]

Can you expand on why you believe this to be a problem? Type-stability for type-stability sake is a non-goal.

[bgctw]

My specific case: I have a computation within a hybrid process-based and machine learning model, where I multiply several vectors by a correctly shaped Unit-Vectors to form higher-dimensional arrays by recycling in the dot notation. The results of these multiplications are used in downstream and it is not practical to implement function barriers. More importantly, also the result of the function (that does this computation) is used in downstream computations and I do not want the user of the function to care for type stability. Hence, a type-stable ones is important to me.

Currently I overload my own function ones_similar_x(x, size) for AbstractArray and cuArray to return FillArrays.Ones(size) or CUDA.ones(size) respectively. This function is type stable.
KernelAbstractions would be a better alternative because many users can use the functionality and it needs to be maintained only for one central package, and extended for other GPU architectures beyond CUDA. I hope for more functionality, e.g. rand, that takes care of the backend.
In such cases where just an Array is created on the GPU, do I need to call synchronize after the creation?