Bump DiffEqFlux in tests

test/Project.toml

@@ -1,10 +1,12 @@
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 DiffEqFlux = "aae7a2af-3d4f-5e19-a356-7da93b79d9d0"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
@@ -19,11 +21,13 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
-DiffEqFlux = ">= 1.43.0"
+ComponentArrays = ">= 0.13.9"
+DiffEqFlux = ">= 2"
 FiniteDiff = ">= 2.8.1"
 Optimisers = ">= 0.2.5"
 ForwardDiff = ">= 0.10.19"
 IterTools = ">= 1.3.0"
+Lux = ">= 0.4.50"
 ModelingToolkit = ">= 8.11.0"
 Optim = ">= 1.4.1"
 OrdinaryDiffEq = ">= 5"

test/diffeqfluxtests.jl

@@ -1,5 +1,6 @@
-using OrdinaryDiffEq, DiffEqFlux, Optimization, OptimizationOptimJL, OptimizationOptimisers,
-      ForwardDiff
+using OrdinaryDiffEq, DiffEqFlux, Lux, Optimization, OptimizationOptimJL,
+      OptimizationOptimisers, ForwardDiff, ComponentArrays, Random
+rng = Random.default_rng()
 
 function lotka_volterra!(du, u, p, t)
     x, y = u
@@ -68,17 +69,15 @@ end
 prob_trueode = ODEProblem(trueODEfunc, u0, tspan)
 ode_data = Array(solve(prob_trueode, Tsit5(), saveat = tsteps))
 
-dudt2 = FastChain((x, p) -> x .^ 3,
-                  FastDense(2, 50, tanh),
-                  FastDense(50, 2))
+dudt2 = Lux.Chain(x -> x .^ 3,
+                  Lux.Dense(2, 50, tanh),
+                  Lux.Dense(50, 2))
 prob_neuralode = NeuralODE(dudt2, tspan, Tsit5(), saveat = tsteps)
-
-dudt2 = Chain(x -> x .^ 3,
-              Dense(2, 50, tanh),
-              Dense(50, 2))
+pp, st = Lux.setup(rng, dudt2)
+pp = ComponentArray(pp)
 
 function predict_neuralode(p)
-    Array(prob_neuralode(u0, p))
+    Array(prob_neuralode(u0, p, st)[1])
 end
 
 function loss_neuralode(p)
@@ -98,7 +97,7 @@ end
 
 optprob = OptimizationFunction((p, x) -> loss_neuralode(p), Optimization.AutoForwardDiff())
 
-prob = Optimization.OptimizationProblem(optprob, prob_neuralode.p)
+prob = Optimization.OptimizationProblem(optprob, pp)
 
 result_neuralode = Optimization.solve(prob,
                                       OptimizationOptimisers.ADAM(), callback = callback,

test/minibatch.jl

@@ -1,5 +1,7 @@
 using DiffEqFlux, Optimization, OrdinaryDiffEq, OptimizationOptimisers, ModelingToolkit,
-      SciMLSensitivity
+      SciMLSensitivity, Lux, Random, ComponentArrays
+
+rng = Random.default_rng()
 
 function newtons_cooling(du, u, p, t)
     temp = u[1]
@@ -13,7 +15,7 @@ function true_sol(du, u, p, t)
 end
 
 function dudt_(u, p, t)
-    ann(u, p) .* u
+    ann(u, p, st)[1] .* u
 end
 
 callback = function (p, l, pred; doplot = false) #callback function to observe training
@@ -35,8 +37,10 @@ t = range(tspan[1], tspan[2], length = datasize)
 true_prob = ODEProblem(true_sol, u0, tspan)
 ode_data = Array(solve(true_prob, Tsit5(), saveat = t))
 
-ann = FastChain(FastDense(1, 8, tanh), FastDense(8, 1, tanh))
-pp = initial_params(ann)
+ann = Lux.Chain(Lux.Dense(1, 8, tanh), Lux.Dense(8, 1, tanh))
+pp, st = Lux.setup(rng, ann)
+pp = ComponentArray(pp)
+
 prob = ODEProblem{false}(dudt_, u0, tspan, pp)
 
 function predict_adjoint(fullp, time_batch)