Merge pull request #935 from SciML/ChrisRackauckas-patch-2-1

Fix mixed GPU CPU adjoint test by switching to Lux

test/gpu/mixed_gpu_cpu_adjoint.jl

@@ -1,14 +1,17 @@
 using SciMLSensitivity, OrdinaryDiffEq
-using Flux, CUDA, Test, Zygote, Random, LinearAlgebra
+using Lux, CUDA, Test, Zygote, Random, LinearAlgebra, ComponentArrays
 
 CUDA.allowscalar(false)
 
 H = CuArray(rand(Float32, 2, 2))
-ann = Chain(Dense(1, 4, tanh))
-p, re = Flux.destructure(ann)
+ann = Lux.Chain(Lux.Dense(1, 4, tanh))
+rng = Random.default_rng()
+p, st = Lux.setup(rng, ann)
+p = ComponentArray(p)
+const _st = st
 
 function func(x, p, t)
-    (re(p)([t])[1] * H) * x
+   CuArray(reshape(first(ann([t],p,_st)),2,2)) * H * x
 end
 
 x0 = CuArray(rand(Float32, 2))
@@ -30,10 +33,7 @@ function cost(p)
 end
 
 grad = Zygote.gradient(cost, p)[1]
-@test !iszero(grad[1])
-@test iszero(grad[2:4])
-@test !iszero(grad[5])
-@test iszero(grad[6:end])
+@test !iszero(grad)
 
 ###
 # https://github.com/SciML/SciMLSensitivity.jl/issues/632