Fix Part 4

Author: SebastianM-C

tutorials/exercises/02-workshop_solutions.jmd

@@ -195,6 +195,7 @@ function henon(dz,z,p,t)
   dq₂ = p₂
 
   dz .= [dp₁, dp₂, dq₁, dq₂]
+  return nothing
 end
 
 u₀ = [0.1, 0.0, 0.0, 0.5]
@@ -240,7 +241,7 @@ function generate_ics(E,n)
   # figures in the Henon-Heiles 1964 article
   qrange = range(-0.4, stop = 1.0, length = n)
   prange = range(-0.5, stop = 0.5, length = n)
-  z0 = Vector{Vector{Float64}}()
+  z0 = Vector{Vector{typeof(E)}}()
   for q in qrange
     V = H([0,0],[0,q],nothing)
     V ≥ E && continue
@@ -269,11 +270,28 @@ plot(sim, vars=(3,4), tspan=(0,10))
 
 ## Part 4: Parallelized GPU Ensemble Solving
 
+In order to use GPU parallelization we must make all inputs
+(initial conditions, tspan, etc.) `Float32` and the function
+definition should be in the in-place form, avoid bound checking and
+return `nothing`.
+
 ```julia
 using DiffEqGPU
 
-z0 = generate_ics(0.125, 100)
-sim = solve(ensprob, Vern9(), EnsembleGPUArray(), trajectories=length(z0))
+function henon_gpu(dz,z,p,t)
+  @inbounds begin
+    dz[1] = -z[3]*(1 + 2z[4])
+    dz[2] = -z[4]-(z[3]^2 - z[4]^2)
+    dz[3] = z[1]
+    dz[4] = z[2]
+  end
+  return nothing
+end
+
+z0 = generate_ics(0.125f0, 50)
+prob_gpu = ODEProblem(henon_gpu, Float32.(u₀), (0.f0, 1000.f0))
+ensprob = EnsembleProblem(prob_gpu, prob_func=prob_func)
+sim = solve(ensprob, Tsit5(), EnsembleGPUArray(), trajectories=length(z0))
 ```
 # Problem 6: Training Neural Stochastic Differential Equations with GPU acceleration (I)