Merge pull request #924 from SciML/ap/updates

Getting tests to pass again

.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: 1 # These tests require quite a lot of GPU memory
+ GROUP: CUDA
  DATADEPS_ALWAYS_ACCEPT: 'true'
  JULIA_PKG_SERVER: "" # it often struggles with our large artifacts
 

.github/workflows/CI.yml

@@ -21,10 +21,10 @@ jobs:
  - BasicNeuralDE
  - AdvancedNeuralDE
  - Newton
- - Aqua
+ - QA
  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

@@ -7,7 +7,6 @@ version = "3.4.0"
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
@@ -23,44 +22,77 @@ RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
+Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [compat]
-ADTypes = "0.2, 1"
-Adapt = "3, 4"
+ADTypes = "1"
+Adapt = "4"
+Aqua = "0.8.7"
+BenchmarkTools = "1.5.0"
+CUDA = "5.3.4"
 ChainRulesCore = "1"
-ComponentArrays = "0.15.5"
+ComponentArrays = "0.15.12"
 ConcreteStructs = "0.2"
+DataInterpolations = "5.0.0"
+DelayDiffEq = "5.47.3"
 DiffEqBase = "6.41"
+DiffEqCallbacks = "3.6.2"
+Distances = "0.10.11"
+Distributed = "1.10"
 Distributions = "0.25"
 DistributionsAD = "0.6"
+ExplicitImports = "1.4.4"
+Flux = "0.14.15"
 ForwardDiff = "0.10"
 Functors = "0.4"
-LinearAlgebra = "<0.0.1, 1"
-Lux = "0.5.5"
+LinearAlgebra = "1.10"
+Lux = "0.5.50"
+LuxCUDA = "0.3.2"
 LuxCore = "0.1"
+MLDataUtils = "0.5.4"
+MLDatasets = "0.7.14"
+NLopt = "1.0.2"
+NNlib = "0.9.16"
+OneHotArrays = "0.2.5"
+Optimisers = "0.3.3"
+Optimization = "3.25.0"
+OptimizationOptimJL = "0.3.0"
+OptimizationOptimisers = "0.2.1"
+OrdinaryDiffEq = "6.76.0"
 PrecompileTools = "1"
-Random = "<0.0.1, 1"
+Printf = "1.10"
+Random = "1.10"
+ReTestItems = "1.24.0"
 RecursiveArrayTools = "2, 3"
 Reexport = "0.2, 1"
+ReverseDiff = "1.15.3"
+SafeTestsets = "0.1.0"
 SciMLBase = "1, 2"
 SciMLSensitivity = "7"
+Setfield = "1.1.1"
+StaticArrays = "1.9.4"
+Statistics = "1.10"
+StochasticDiffEq = "6.65.1"
+Test = "1.10"
 Tracker = "0.2.29"
 Zygote = "0.6"
 ZygoteRules = "0.2"
-julia = "1.9"
+julia = "1.10"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
 DelayDiffEq = "bcd4f6db-9728-5f36-b5f7-82caef46ccdb"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
+ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 MLDataUtils = "cc2ba9b6-d476-5e6d-8eaf-a92d5412d41d"
@@ -74,6 +106,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 +116,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", "ComponentArrays", "DataInterpolations", "DelayDiffEq", "DiffEqCallbacks", "Distances", "Distributed", "ExplicitImports", "Flux", "LuxCUDA", "MLDataUtils", "MLDatasets", "NLopt", "NNlib", "OneHotArrays", "Optimisers", "Optimization", "OptimizationOptimJL", "OptimizationOptimisers", "OrdinaryDiffEq", "Printf", "ReTestItems", "RecursiveArrayTools", "ReverseDiff", "SafeTestsets", "StaticArrays", "Statistics", "StochasticDiffEq", "Test"]

README.md

@@ -16,7 +16,7 @@ by helping users put diffeq solvers into neural networks. This package utilizes
 [Scientific Machine Learning](https://www.stochasticlifestyle.com/the-essential-tools-of-scientific-machine-learning-scientific-ml/), specifically neural differential equations to add physical information into traditional machine learning.
 
 > [!NOTE]
-> We maintain backwards compatibility with [Flux.jl](https://docs.sciml.ai/Flux/stable/) via [Lux.transform](https://lux.csail.mit.edu/stable/api/Lux/flux_to_lux#Lux.transform)
+> We maintain backwards compatibility with [Flux.jl](https://docs.sciml.ai/Flux/stable/) via [FromFluxAdaptor()](https://lux.csail.mit.edu/stable/api/Lux/interop#Lux.FromFluxAdaptor)
 
 ## Tutorials and Documentation
 
@@ -63,7 +63,7 @@ explore various ways to integrate the two methodologies:
 
 ## Breaking Changes in v3
 
- - Flux dependency is dropped. If a non Lux `AbstractExplicitLayer` is passed we try to automatically convert it to a Lux model with `Lux.transform(model)`.
+ - Flux dependency is dropped. If a non Lux `AbstractExplicitLayer` is passed we try to automatically convert it to a Lux model with `FromFluxAdaptor()(model)`.
  - `Flux` is no longer re-exported from `DiffEqFlux`. Instead we reexport `Lux`.
  - `NeuralDAE` now allows an optional `du0` as input.
  - `TensorLayer` is now a Lux Neural Network.

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,6 +1,8 @@
+#! format: off
 pages = [
  "DiffEqFlux.jl: High Level Scientific Machine Learning (SciML) Pre-Built Architectures" => "index.md",
- "Differential Equation Machine Learning Tutorials" => Any["examples/neural_ode.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",
@@ -11,13 +13,19 @@ pages = [
  "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",
+ "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"]
+ "Hamiltonian Neural Network Layer" => "layers/HamiltonianNN.md"
+ ],
+ "Utility Function APIs" => Any[
+ "Smoothed Collocation" => "utilities/Collocation.md",
+ "Multiple Shooting Functionality" => "utilities/MultipleShooting.md"
+ ]
 ]
+#! format: on

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