Diagonal(::CuVector) .* ::Number returns full CuMatrix

[vpuri3]

julia> Diagonal(CUDA.ones(4)) .* 4f0
4×4 CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}:
 4.0  0.0  0.0  0.0
 0.0  4.0  0.0  0.0
 0.0  0.0  4.0  0.0
 0.0  0.0  0.0  4.0

Should return Diagonal{T, CuVector} like

julia> Diagonal(ones(4)) .* 4f0
4×4 Diagonal{Float64, Vector{Float64}}:
 4.0   ⋅    ⋅    ⋅ 
  ⋅   4.0   ⋅    ⋅ 
  ⋅    ⋅   4.0   ⋅ 
  ⋅    ⋅    ⋅   4.0

[vpuri3]

@maleadt i can add the required method if you can LMK where that needs to be done

[maleadt]

I don't know that by heart; you can check how Base's Broadcast implementation preserves the Diagonal and apply the same logic to the GPUArrays broadcast implementation.