Error when running DiffEqGPU example from documentation

Hello,

I ran the following example, copy-pasted from here, and received an UndefVarError.

Here is the example snippet from the linked website:

using DiffEqGPU, OrdinaryDiffEq
function lorenz(du, u, p, t)
    du[1] = p[1] * (u[2] - u[1])
    du[2] = u[1] * (p[2] - u[3]) - u[2]
    du[3] = u[1] * u[2] - p[3] * u[3]
end

u0 = Float32[1.0;0.0;0.0]
tspan = (0.0f0,100.0f0)
p = [10.0f0,28.0f0,8/3f0]
prob = ODEProblem(lorenz,u0,tspan,p)
prob_func = (prob,i,repeat) -> remake(prob,p=rand(Float32,3).*p)
monteprob = EnsembleProblem(prob, prob_func = prob_func, safetycopy=false)
@time sol = solve(monteprob,Tsit5(),EnsembleGPUArray(CUDADevice()),trajectories=10_000,saveat=1.0f0)

Here is the command I use to run the code and the corresponding output:

$ srun -G 4 --pty .local/bin/julia ./research/diffeq/sandbox/www_copypasta_ensemble.jl 
ERROR: LoadError: UndefVarError: `CUDADevice` not defined
Stacktrace:
 [1] top-level scope
   @ ./timing.jl:273
in expression starting at /common/home/jag619/research/diffeq/sandbox/www_copypasta_ensemble.jl:14
srun: error: ilab1: task 0: Exited with exit code 1

Hopefully the following suffices for version information:

julia> 

(@v1.9) pkg> status
Status `~/.julia/environments/v1.9/Project.toml`
  [052768ef] CUDA v4.3.0
  [071ae1c0] DiffEqGPU v2.2.1
  [1dea7af3] OrdinaryDiffEq v6.51.2
  [37e2e46d] LinearAlgebra

julia> using CUDA

julia> CUDA.versioninfo()
CUDA runtime 12.1, artifact installation
CUDA driver 12.1
NVIDIA driver 470.161.3, originally for CUDA 11.4

CUDA libraries: 
- CUBLAS: 12.1.3
- CURAND: 10.3.2
- CUFFT: 11.0.2
- CUSOLVER: 11.4.5
- CUSPARSE: 12.1.0
- CUPTI: 18.0.0
- NVML: 11.0.0+470.161.3

Julia packages: 
- CUDA.jl: 4.3.0
- CUDA_Driver_jll: 0.5.0+1
- CUDA_Runtime_jll: 0.6.0+0
- CUDA_Runtime_Discovery: 0.2.2

Toolchain:
- Julia: 1.9.0
- LLVM: 14.0.6
- PTX ISA support: 3.2, 4.0, 4.1, 4.2, 4.3, 5.0, 6.0, 6.1, 6.3, 6.4, 6.5, 7.0, 7.1, 7.2, 7.3, 7.4, 7.5
- Device capability support: sm_37, sm_50, sm_52, sm_53, sm_60, sm_61, sm_62, sm_70, sm_72, sm_75, sm_80, sm_86

4 devices:
  0: NVIDIA RTX A4000 (sm_86, 14.742 GiB / 14.746 GiB available)
  1: NVIDIA RTX A4000 (sm_86, 14.742 GiB / 14.746 GiB available)
  2: NVIDIA RTX A4000 (sm_86, 14.742 GiB / 14.746 GiB available)
  3: NVIDIA RTX A4000 (sm_86, 14.742 GiB / 14.746 GiB available)

Am I missing something here? Where is CUDADevice usually imported from?

This is a KernelAbstractions.jl concept. It’s been renamed to CUDABackend, and is now part of CUDA.jl itself.

Sorry about that, the docstrings are fixed now via Fix docstring examples by adding `using CUDA` by ChrisRackauckas · Pull Request #279 · SciML/DiffEqGPU.jl · GitHub. You just need a using CUDA in there to specify which GPU package you’re using.

using DiffEqGPU, CUDA, OrdinaryDiffEq
function lorenz(du, u, p, t)
    du[1] = p[1] * (u[2] - u[1])
    du[2] = u[1] * (p[2] - u[3]) - u[2]
    du[3] = u[1] * u[2] - p[3] * u[3]
end

u0 = Float32[1.0;0.0;0.0]
tspan = (0.0f0,100.0f0)
p = [10.0f0,28.0f0,8/3f0]
prob = ODEProblem(lorenz,u0,tspan,p)
prob_func = (prob,i,repeat) -> remake(prob,p=rand(Float32,3).*p)
monteprob = EnsembleProblem(prob, prob_func = prob_func, safetycopy=false)
@time sol = solve(monteprob,Tsit5(),EnsembleGPUArray(CUDADevice()),trajectories=10_000,saveat=1.0f0)