BFGS Errors with GPU

[collinwarner]

I am running into an error with using BFGS and cu arrays I have created an MWE:

using DifferentialEquations, Flux, Optim, DiffEqFlux, CuArrays
function lotka_volterra(du,u,p,t)
  x, y = u
  α, β, δ, γ = p
  du[1] = dx = α*x - β*x*y
  du[2] = dy = -δ*y + γ*x*y
end
u0 = [1.0,1.0] |>gpu
tspan = (0.0,10.0)
p = [1.5,1.0,3.0,1.0] |>gpu
prob = ODEProblem(lotka_volterra,u0,tspan,p)
sol = solve(prob,Tsit5())
function predict_adjoint(p) 
    Array(concrete_solve(prob,Tsit5(),u0,p,saveat=0.0:0.1:10.0))
end

function loss_adjoint(p)
    prediction = predict_adjoint(p)
    loss = sum(abs2,x-1 for x in prediction)
    loss,prediction
end

cb = function (p,l,pred) #callback function to observe training
    display(l)
    false
end
  
cb(p,loss_adjoint(p)...)
  
res = DiffEqFlux.sciml_train(loss_adjoint, p, BFGS(initial_stepnorm = 0.0001), cb = cb)

The error looks like it occurs when initial_invH is intialized in BFGS.jl, since it is created as an array and not a cuArray which creates a multiplication error. The error message is below.:

> ┌ Warning: Performing scalar operations on GPU arrays: This is very slow, consider disallowing these operations with `allowscalar(false)`
└ @ GPUArrays C:\Users\Collin\.julia\packages\GPUArrays\WZupy\src\host\indexing.jl:43
1015.00555f0
1015.00555f0
ERROR: LoadError: ArgumentError: cannot take the CPU address of a CuArray{Float32,1,Nothing}
Stacktrace:
 [1] unsafe_convert(::Type{Ptr{Float32}}, ::CuArray{Float32,1,Nothing}) at C:\Users\Collin\.julia\packages\CuArrays\4Q1BY\src\array.jl:226
 [2] gemv!(::Char, ::Float32, ::Array{Float32,2}, ::CuArray{Float32,1,Nothing}, ::Float32, ::CuArray{Float32,1,Nothing}) at D:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.4\LinearAlgebra\src\blas.jl:587
 [3] gemv!(::CuArray{Float32,1,Nothing}, ::Char, ::Array{Float32,2}, ::CuArray{Float32,1,Nothing}, ::Bool, ::Bool) at D:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.4\LinearAlgebra\src\matmul.jl:470
 [4] mul! at D:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.4\LinearAlgebra\src\matmul.jl:66 [inlined]
 [5] mul! at D:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.4\LinearAlgebra\src\matmul.jl:208 [inlined]
 [6] update_state!(::TwiceDifferentiable{Float32,CuArray{Float32,1,Nothing},Array{Float32,2},CuArray{Float32,1,Nothing}}, ::Optim.BFGSState{CuArray{Float32,1,Nothing},Array{Float32,2},Float32,CuArray{Float32,1,Nothing}}, ::BFGS{LineSearches.InitialStatic{Float64},LineSearches.HagerZhang{Float64,Base.RefValue{Bool}},Nothing,Float64,Flat}) at C:\Users\Collin\.julia\packages\Optim\UkDyx\src\multivariate\solvers\first_order\bfgs.jl:100
 [7] optimize(::TwiceDifferentiable{Float32,CuArray{Float32,1,Nothing},Array{Float32,2},CuArray{Float32,1,Nothing}}, ::CuArray{Float32,1,Nothing}, ::BFGS{LineSearches.InitialStatic{Float64},LineSearches.HagerZhang{Float64,Base.RefValue{Bool}},Nothing,Float64,Flat}, ::Optim.Options{Float64,DiffEqFlux.var"#_cb#46"{var"#22#23",Base.Iterators.Cycle{Tuple{DiffEqFlux.NullData}}}}, ::Optim.BFGSState{CuArray{Float32,1,Nothing},Array{Float32,2},Float32,CuArray{Float32,1,Nothing}}) at C:\Users\Collin\.julia\packages\Optim\UkDyx\src\multivariate\optimize\optimize.jl:57 [8] optimize(::TwiceDifferentiable{Float32,CuArray{Float32,1,Nothing},Array{Float32,2},CuArray{Float32,1,Nothing}}, ::CuArray{Float32,1,Nothing}, ::BFGS{LineSearches.InitialStatic{Float64},LineSearches.HagerZhang{Float64,Base.RefValue{Bool}},Nothing,Float64,Flat}, ::Optim.Options{Float64,DiffEqFlux.var"#_cb#46"{var"#22#23",Base.Iterators.Cycle{Tuple{DiffEqFlux.NullData}}}}) at C:\Users\Collin\.julia\packages\Optim\UkDyx\src\multivariate\optimize\optimize.jl:33
 [9] sciml_train(::Function, ::CuArray{Float32,1,Nothing}, ::BFGS{LineSearches.InitialStatic{Float64},LineSearches.HagerZhang{Float64,Base.RefValue{Bool}},Nothing,Float64,Flat}, ::Base.Iterators.Cycle{Tuple{DiffEqFlux.NullData}}; cb::Function, maxiters::Int64, diffmode::DiffEqFlux.ZygoteDiffMode, kwargs::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}})
at C:\Users\Collin\.julia\packages\DiffEqFlux\gMXy7\src\train.jl:269
 [10] top-level scope at C:\Users\Collin\Documents\Downloads\universal_differential_equations-1-master\universal_differential_equations-1-master\Climate\NeuralPDE\MWE.jl:32
 [11] include(::String) at .\client.jl:439
 [12] top-level scope at REPL[3]:1
in expression starting at C:\Users\Collin\Documents\Downloads\universal_differential_equations-1-master\universal_differential_equations-1-master\Climate\NeuralPDE\MWE.jl:32

[ChrisRackauckas]

@pkofod are there plans to support CuArrays in Optim?

[pkofod]

Ah this is because we loop for the update?

[ChrisRackauckas]

It looks like it might be a 5 argument mul! issue.

[ChrisRackauckas]

Fixed upstream.