Clean up the tests

Author: avik-pal

.JuliaFormatter.toml

@@ -2,4 +2,5 @@ style = "sciml"
 annotate_untyped_fields_with_any = false
 format_markdown = true
 format_docstrings = true
-separate_kwargs_with_semicolon = true
+separate_kwargs_with_semicolon = true
+join_lines_based_on_source = false

.buildkite/pipeline.yml

@@ -12,7 +12,8 @@ steps:
     # Don't run Buildkite if the commit message includes the text [skip tests]
     if: build.message !~ /\[skip tests\]/
     env:
-      GROUP: GPU
+      RETESTITEMS_NWORKERS: 4
+      GROUP: CUDA
       DATADEPS_ALWAYS_ACCEPT: 'true'
       JULIA_PKG_SERVER: "" # it often struggles with our large artifacts
 

.github/workflows/CI.yml

@@ -24,7 +24,7 @@ jobs:
           - Aqua
         version:
           - '1'
-          - '~1.10.0-0'
+          - '1.10'
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2
@@ -46,6 +46,7 @@ jobs:
           coverage: false
         env:
           GROUP: ${{ matrix.group }}
+          RETESTITEMS_NWORKERS: 4
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v4
         with:

.github/workflows/FormatCheck.yml

@@ -1,42 +1,9 @@
-name: format-check
+name: Format suggestions
 
-on:
-  push:
-    branches:
-      - 'master'
-      - 'release-'
-    tags: '*'
-  pull_request:
+on: [pull_request]
 
 jobs:
-  build:
-    runs-on: ${{ matrix.os }}
-    strategy:
-      matrix:
-        julia-version: [1]
-        julia-arch: [x86]
-        os: [ubuntu-latest]
+  code-style:
+    runs-on: ubuntu-latest
     steps:
-      - uses: julia-actions/setup-julia@latest
-        with:
-          version: ${{ matrix.julia-version }}
-
-      - uses: actions/checkout@v4
-      - name: Install JuliaFormatter and format
-        # This will use the latest version by default but you can set the version like so:
-        #
-        # julia  -e 'using Pkg; Pkg.add(PackageSpec(name="JuliaFormatter", version="0.13.0"))'
-        run: |
-          julia  -e 'using Pkg; Pkg.add(PackageSpec(name="JuliaFormatter"))'
-          julia  -e 'using JuliaFormatter; format(".", verbose=true)'
-      - name: Format check
-        run: |
-          julia -e '
-          out = Cmd(`git diff --name-only`) |> read |> String
-          if out == ""
-              exit(0)
-          else
-              @error "Some files have not been formatted !!!"
-              write(stdout, out)
-              exit(1)
-          end'
+      - uses: julia-actions/julia-format@v2

Project.toml

@@ -28,8 +28,8 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [compat]
-ADTypes = "0.2, 1"
-Adapt = "3, 4"
+ADTypes = "1"
+Adapt = "4"
 ChainRulesCore = "1"
 ComponentArrays = "0.15.5"
 ConcreteStructs = "0.2"
@@ -38,19 +38,19 @@ Distributions = "0.25"
 DistributionsAD = "0.6"
 ForwardDiff = "0.10"
 Functors = "0.4"
-LinearAlgebra = "<0.0.1, 1"
-Lux = "0.5.5"
+LinearAlgebra = "1.10"
+Lux = "0.5.50"
 LuxCore = "0.1"
 PrecompileTools = "1"
-Random = "<0.0.1, 1"
+Random = "1.10"
 RecursiveArrayTools = "2, 3"
 Reexport = "0.2, 1"
 SciMLBase = "1, 2"
 SciMLSensitivity = "7"
 Tracker = "0.2.29"
 Zygote = "0.6"
 ZygoteRules = "0.2"
-julia = "1.9"
+julia = "1.10"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
@@ -74,6 +74,7 @@ OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e"
 OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
@@ -83,4 +84,4 @@ StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "BenchmarkTools", "CUDA", "DataInterpolations", "DelayDiffEq", "DiffEqCallbacks", "Distances", "Distributed", "Flux", "LuxCUDA", "MLDataUtils", "MLDatasets", "NLopt", "NNlib", "OneHotArrays", "Optimisers", "Optimization",  "OptimizationOptimJL",  "OptimizationOptimisers", "OrdinaryDiffEq", "Printf", "RecursiveArrayTools", "ReverseDiff", "SafeTestsets", "StaticArrays", "Statistics", "StochasticDiffEq", "Test"]
+test = ["Aqua", "BenchmarkTools", "CUDA", "DataInterpolations", "DelayDiffEq", "DiffEqCallbacks", "Distances", "Distributed", "Flux", "LuxCUDA", "MLDataUtils", "MLDatasets", "NLopt", "NNlib", "OneHotArrays", "Optimisers", "Optimization", "OptimizationOptimJL", "OptimizationOptimisers", "OrdinaryDiffEq", "Printf", "ReTestItems", "RecursiveArrayTools", "ReverseDiff", "SafeTestsets", "StaticArrays", "Statistics", "StochasticDiffEq", "Test"]

docs/make.jl

@@ -11,12 +11,13 @@ include("pages.jl")
 
 makedocs(; sitename = "DiffEqFlux.jl",
     authors = "Chris Rackauckas et al.",
-    clean = true, doctest = false, linkcheck = true,
+    clean = true,
+    doctest = false,
+    linkcheck = true,
     warnonly = [:docs_block, :missing_docs],
     modules = [DiffEqFlux],
     format = Documenter.HTML(; assets = ["assets/favicon.ico"],
         canonical = "https://docs.sciml.ai/DiffEqFlux/stable/"),
     pages = pages)
 
-deploydocs(; repo = "github.com/SciML/DiffEqFlux.jl.git",
-    push_preview = true)
+deploydocs(; repo = "github.com/SciML/DiffEqFlux.jl.git", push_preview = true)

docs/pages.jl

@@ -1,23 +1,17 @@
 pages = [
     "DiffEqFlux.jl: High Level Scientific Machine Learning (SciML) Pre-Built Architectures" => "index.md",
-    "Differential Equation Machine Learning Tutorials" => Any["examples/neural_ode.md",
-        "examples/GPUs.md",
-        "examples/mnist_neural_ode.md",
-        "examples/mnist_conv_neural_ode.md",
-        "examples/augmented_neural_ode.md",
-        "examples/neural_sde.md",
-        "examples/collocation.md",
-        "examples/normalizing_flows.md",
-        "examples/hamiltonian_nn.md",
-        "examples/tensor_layer.md",
-        "examples/multiple_shooting.md",
-        "examples//neural_ode_weather_forecast.md"],
+    "Differential Equation Machine Learning Tutorials" => Any[
+        "examples/neural_ode.md", "examples/GPUs.md",
+        "examples/mnist_neural_ode.md", "examples/mnist_conv_neural_ode.md",
+        "examples/augmented_neural_ode.md", "examples/neural_sde.md",
+        "examples/collocation.md", "examples/normalizing_flows.md",
+        "examples/hamiltonian_nn.md", "examples/tensor_layer.md",
+        "examples/multiple_shooting.md", "examples//neural_ode_weather_forecast.md"],
     "Layer APIs" => Any["Classical Basis Layers" => "layers/BasisLayers.md",
         "Tensor Product Layer" => "layers/TensorLayer.md",
         "Continuous Normalizing Flows Layer" => "layers/CNFLayer.md",
         "Spline Layer" => "layers/SplineLayer.md",
         "Neural Differential Equation Layers" => "layers/NeuralDELayers.md",
         "Hamiltonian Neural Network Layer" => "layers/HamiltonianNN.md"],
     "Utility Function APIs" => Any["Smoothed Collocation" => "utilities/Collocation.md",
-        "Multiple Shooting Functionality" => "utilities/MultipleShooting.md"]
-]
+        "Multiple Shooting Functionality" => "utilities/MultipleShooting.md"]]

docs/src/examples/GPUs.md

@@ -9,7 +9,7 @@ For a detailed discussion on how GPUs need to be setup refer to
 
 ```@example gpu
 using OrdinaryDiffEq, Lux, LuxCUDA, SciMLSensitivity, ComponentArrays, Random
-rng = Random.default_rng()
+rng = Xoshiro(0)
 
 const cdev = cpu_device()
 const gdev = gpu_device()
@@ -72,14 +72,14 @@ Here is the full neural ODE example. Note that we use the `gpu_device` function
 same code works on CPUs and GPUs, dependent on `using LuxCUDA`.
 
 ```@example gpu
-using Lux, Optimization, OptimizationOptimisers, Zygote, OrdinaryDiffEq,
-      Plots, LuxCUDA, SciMLSensitivity, Random, ComponentArrays
+using Lux, Optimization, OptimizationOptimisers, Zygote, OrdinaryDiffEq, Plots, LuxCUDA,
+      SciMLSensitivity, Random, ComponentArrays
 import DiffEqFlux: NeuralODE
 
 CUDA.allowscalar(false) # Makes sure no slow operations are occurring
 
 #rng for Lux.setup
-rng = Random.default_rng()
+rng = Xoshiro(0)
 # Generate Data
 u0 = Float32[2.0; 0.0]
 datasize = 30

docs/src/examples/augmented_neural_ode.md

@@ -42,15 +42,16 @@ function construct_model(out_dim, input_dim, hidden_dim, augment_dim)
     input_dim = input_dim + augment_dim
     node = NeuralODE(
         Chain(Dense(input_dim, hidden_dim, relu),
-            Dense(hidden_dim, hidden_dim, relu),
-            Dense(hidden_dim, input_dim)),
+            Dense(hidden_dim, hidden_dim, relu), Dense(hidden_dim, input_dim)),
         (0.0f0, 1.0f0),
         Tsit5();
         save_everystep = false,
-        reltol = 1.0f-3, abstol = 1.0f-3, save_start = false)
+        reltol = 1.0f-3,
+        abstol = 1.0f-3,
+        save_start = false)
     node = augment_dim == 0 ? node : AugmentedNDELayer(node, augment_dim)
     model = Chain(node, diffeqarray_to_array, Dense(input_dim, out_dim))
-    ps, st = Lux.setup(Random.default_rng(), model)
+    ps, st = Lux.setup(Xoshiro(0), model)
     return model, ps |> gdev, st |> gdev
 end
 
@@ -70,8 +71,8 @@ end
 
 loss_node(model, x, y, ps, st) = mean((first(model(x, ps, st)) .- y) .^ 2)
 
-dataloader = concentric_sphere(2, (0.0f0, 2.0f0), (3.0f0, 4.0f0), 2000, 2000;
-    batch_size = 256)
+dataloader = concentric_sphere(
+    2, (0.0f0, 2.0f0), (3.0f0, 4.0f0), 2000, 2000; batch_size = 256)
 
 iter = 0
 cb = function (ps, l)
@@ -186,15 +187,16 @@ function construct_model(out_dim, input_dim, hidden_dim, augment_dim)
     input_dim = input_dim + augment_dim
     node = NeuralODE(
         Chain(Dense(input_dim, hidden_dim, relu),
-            Dense(hidden_dim, hidden_dim, relu),
-            Dense(hidden_dim, input_dim)),
+            Dense(hidden_dim, hidden_dim, relu), Dense(hidden_dim, input_dim)),
         (0.0f0, 1.0f0),
         Tsit5();
         save_everystep = false,
-        reltol = 1.0f-3, abstol = 1.0f-3, save_start = false)
+        reltol = 1.0f-3,
+        abstol = 1.0f-3,
+        save_start = false)
     node = augment_dim == 0 ? node : AugmentedNDELayer(node, augment_dim)
     model = Chain(node, diffeqarray_to_array, Dense(input_dim, out_dim))
-    ps, st = Lux.setup(Random.default_rng(), model)
+    ps, st = Lux.setup(Xoshiro(0), model)
     return model, ps |> gdev, st |> gdev
 end
 ```
@@ -236,8 +238,8 @@ Next, we generate the dataset. We restrict ourselves to 2 dimensions as it is ea
 We sample a total of `4000` data points.
 
 ```@example augneuralode
-dataloader = concentric_sphere(2, (0.0f0, 2.0f0), (3.0f0, 4.0f0), 2000, 2000;
-    batch_size = 256)
+dataloader = concentric_sphere(
+    2, (0.0f0, 2.0f0), (3.0f0, 4.0f0), 2000, 2000; batch_size = 256)
 ```
 
 #### Callback Function

docs/src/examples/collocation.md

@@ -19,7 +19,7 @@ using ComponentArrays, Lux, DiffEqFlux, OrdinaryDiffEq, SciMLSensitivity, Optimi
       OptimizationOptimisers, Plots
 
 using Random
-rng = Random.default_rng()
+rng = Xoshiro(0)
 
 u0 = Float32[2.0; 0.0]
 datasize = 300
@@ -101,11 +101,11 @@ The smoothed collocation is a spline fit of the data points which allows
 us to get an estimate of the approximate noiseless dynamics:
 
 ```@example collocation
-using ComponentArrays,
-      Lux, DiffEqFlux, Optimization, OptimizationOptimisers, OrdinaryDiffEq, Plots
+using ComponentArrays, Lux, DiffEqFlux, Optimization, OptimizationOptimisers,
+      OrdinaryDiffEq, Plots
 
 using Random
-rng = Random.default_rng()
+rng = Xoshiro(0)
 
 u0 = Float32[2.0; 0.0]
 datasize = 300

docs/src/examples/hamiltonian_nn.md

@@ -34,7 +34,7 @@ dataloader = ncycle(
     NEPOCHS)
 
 hnn = HamiltonianNN(Chain(Dense(2 => 64, relu), Dense(64 => 1)); ad = AutoZygote())
-ps, st = Lux.setup(Random.default_rng(), hnn)
+ps, st = Lux.setup(Xoshiro(0), hnn)
 ps_c = ps |> ComponentArray
 
 opt = OptimizationOptimisers.Adam(0.01f0)
@@ -57,8 +57,8 @@ res = Optimization.solve(opt_prob, opt, dataloader; callback)
 
 ps_trained = res.u
 
-model = NeuralHamiltonianDE(hnn, (0.0f0, 1.0f0), Tsit5(); save_everystep = false,
-    save_start = true, saveat = t)
+model = NeuralHamiltonianDE(
+    hnn, (0.0f0, 1.0f0), Tsit5(); save_everystep = false, save_start = true, saveat = t)
 
 pred = Array(first(model(data[:, 1], ps_trained, st)))
 plot(data[1, :], data[2, :]; lw = 4, label = "Original")
@@ -101,7 +101,7 @@ We parameterize the HamiltonianNN with a small MultiLayered Perceptron. HNNs are
 
 ```@example hamiltonian
 hnn = HamiltonianNN(Chain(Dense(2 => 64, relu), Dense(64 => 1)); ad = AutoZygote())
-ps, st = Lux.setup(Random.default_rng(), hnn)
+ps, st = Lux.setup(Xoshiro(0), hnn)
 ps_c = ps |> ComponentArray
 
 opt = OptimizationOptimisers.Adam(0.01f0)
@@ -116,8 +116,8 @@ function callback(ps, loss, pred)
     return false
 end
 
-opt_func = OptimizationFunction((ps, _, data, target) -> loss_function(ps, data, target),
-    Optimization.AutoZygote())
+opt_func = OptimizationFunction(
+    (ps, _, data, target) -> loss_function(ps, data, target), Optimization.AutoZygote())
 opt_prob = OptimizationProblem(opt_func, ps_c)
 
 res = solve(opt_prob, opt, dataloader; callback)
@@ -130,8 +130,8 @@ ps_trained = res.u
 In order to visualize the learned trajectories, we need to solve the ODE. We will use the `NeuralHamiltonianDE` layer, which is essentially a wrapper over `HamiltonianNN` layer, and solves the ODE.
 
 ```@example hamiltonian
-model = NeuralHamiltonianDE(hnn, (0.0f0, 1.0f0), Tsit5(); save_everystep = false,
-    save_start = true, saveat = t)
+model = NeuralHamiltonianDE(
+    hnn, (0.0f0, 1.0f0), Tsit5(); save_everystep = false, save_start = true, saveat = t)
 
 pred = Array(first(model(data[:, 1], ps_trained, st)))
 plot(data[1, :], data[2, :]; lw = 4, label = "Original")

docs/src/examples/mnist_conv_neural_ode.md

@@ -7,9 +7,8 @@ For a step-by-step tutorial see the tutorial on the MNIST Neural ODE Classificat
 using Fully Connected Layers.
 
 ```@example mnist_cnn
-using DiffEqFlux, Statistics,
-      ComponentArrays, CUDA, Zygote, MLDatasets, OrdinaryDiffEq, Printf, Test, LuxCUDA,
-      Random
+using DiffEqFlux, Statistics, ComponentArrays, CUDA, Zygote, MLDatasets, OrdinaryDiffEq,
+      Printf, Test, LuxCUDA, Random
 using Optimization, OptimizationOptimisers
 using MLDatasets: MNIST
 using MLDataUtils: LabelEnc, convertlabel, stratifiedobs, batchview
@@ -45,8 +44,8 @@ const bs = 128
 x_train, y_train = loadmnist(bs)
 
 down = Chain(Conv((3, 3), 1 => 64, relu; stride = 1), GroupNorm(64, 64),
-    Conv((4, 4), 64 => 64, relu; stride = 2, pad = 1), GroupNorm(64, 64),
-    Conv((4, 4), 64 => 64; stride = 2, pad = 1))
+    Conv((4, 4), 64 => 64, relu; stride = 2, pad = 1),
+    GroupNorm(64, 64), Conv((4, 4), 64 => 64; stride = 2, pad = 1))
 
 dudt = Chain(Conv((3, 3), 64 => 64, tanh; stride = 1, pad = 1),
     Conv((3, 3), 64 => 64, tanh; stride = 1, pad = 1))
@@ -66,7 +65,7 @@ m = Chain(down,                 # (28, 28, 1, BS) -> (6, 6, 64, BS)
     nn_ode,               # (6, 6, 64, BS) -> (6, 6, 64, BS, 1)
     DiffEqArray_to_Array, # (6, 6, 64, BS, 1) -> (6, 6, 64, BS)
     fc)                   # (6, 6, 64, BS) -> (10, BS)
-ps, st = Lux.setup(Random.default_rng(), m)
+ps, st = Lux.setup(Xoshiro(0), m)
 ps = ComponentArray(ps) |> gdev
 st = st |> gdev
 
@@ -105,8 +104,8 @@ loss_function(ps, x_train[1], y_train[1])
 opt = OptimizationOptimisers.Adam(0.05)
 iter = 0
 
-opt_func = OptimizationFunction((ps, _, x, y) -> loss_function(ps, x, y),
-    Optimization.AutoZygote())
+opt_func = OptimizationFunction(
+    (ps, _, x, y) -> loss_function(ps, x, y), Optimization.AutoZygote())
 opt_prob = OptimizationProblem(opt_func, ps)
 
 function callback(ps, l, pred)