Can't calculate `von_neumann_entropy()` in gpu

[zipeilee]

When calculating the von_neumann_entropy() of a certain quantum state on the GPU, I often receive CUDA errors like

ERROR: ArgumentError: cannot take the CPU address of a CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}
Stacktrace:
 [1] unsafe_convert(#unused#::Type{Ptr{Float32}}, x::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
   @ CUDA ~/.julia/packages/CUDA/pCcGc/src/array.jl:370
 [2] geevx!(balanc::Char, jobvl::Char, jobvr::Char, sense::Char, A::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
   @ LinearAlgebra.LAPACK ~/packages/julias/julia-1.9/share/julia/stdlib/v1.9/LinearAlgebra/src/lapack.jl:2093
 [3] eigvals!(A::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}; permute::Bool, scale::Bool, sortby::typeof(LinearAlgebra.eigsortby))
   @ LinearAlgebra ~/packages/julias/julia-1.9/share/julia/stdlib/v1.9/LinearAlgebra/src/eigen.jl:306
 [4] eigvals!(A::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
   @ LinearAlgebra ~/packages/julias/julia-1.9/share/julia/stdlib/v1.9/LinearAlgebra/src/eigen.jl:304
 [5] eigvals(A::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}; kws::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
   @ LinearAlgebra ~/packages/julias/julia-1.9/share/julia/stdlib/v1.9/LinearAlgebra/src/eigen.jl:339
 [6] eigvals(A::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
   @ LinearAlgebra ~/packages/julias/julia-1.9/share/julia/stdlib/v1.9/LinearAlgebra/src/eigen.jl:339

Then I remind that function von_neumann_entropy() have a operattion p = max.(eigvals(dm), eps(real(eltype(dm)))) but it seems eigvals() operation have not been implemented in CUDA.jl.

I wander if there is any way to improve this operation in Yao.jl so that it can run successfully on the GPU, or alternatively, to only load the quantum state onto the CPU when calculating the entropy?

[zipeilee]

I have tried change the
p = max.(eigvals(dm), eps(real(eltype(dm)))) to

 _, l = eigen(dm)
p = max.(l, eps(real(eltype(dm))))

but due to computational precision, the density matrix is not a Hermitian matrix and LinearAlgebra.eigen() only seems to work for symmetric/hermitian inputs.

[GiggleLiu]

Thank you for the report. I think you can dispatch the method to CPU easily with the following code. I am not very confident that dispatch GPU code to that on CPU is a good default behavior.

julia> using Yao, CUDA

julia> reg = rand_state(6)
ArrayReg{2, ComplexF64, Array...}
    active qubits: 6/6
    nlevel: 2

julia> Yao.von_neumann_entropy(reg::ArrayReg{D, T, <:CuArray{T}}, locs) where {T, D} = von_neumann_entropy(cpu(reg), locs)

julia> von_neumann_entropy(cu(reg), 2:4)
1.659399675203455

Also, if you want to tell the compiler that your matrix is hermitian, please use the LinearAlgebra.Hermitian wrapper.

julia> A = randn(ComplexF64, 4, 4)
4×4 Matrix{ComplexF64}:
 -0.545783+0.477621im   1.36153-0.602449im   0.319406+0.560202im   -0.667118+1.15637im
 -0.203733-0.303557im   1.29376+0.723739im  -0.207717+0.811873im    0.416798-1.11345im
  0.637101+0.968772im  0.950193+1.03022im   -0.048993-0.784866im     1.01995+1.46875im
 -0.929421+1.86055im   0.480936-0.343407im  -0.201151+0.0899719im   0.850967+0.967622im

julia> using LinearAlgebra

julia> eigen(Hermitian(A))
Eigen{ComplexF64, Float64, Matrix{ComplexF64}, Vector{Float64}}
values:
4-element Vector{Float64}:
 -2.6064576696617032
 -1.0277071601685992
  1.933628962937959
  3.2504829331706553
vectors:
4×4 Matrix{ComplexF64}:
  0.246749-0.513774im  -0.394964+0.503586im   -0.299767-0.132123im    -0.349073+0.190754im
 -0.296981+0.345221im  0.0656888-0.0350235im  -0.738583-0.00796948im  -0.377532-0.314618im
 -0.276691-0.271599im  -0.333437-0.655966im   -0.278456-0.0177148im    0.132953+0.461168im
  0.563429-0.0im        0.208308-0.0im        -0.518906-0.0im          0.608188+0.0im