fix: minor corrections

.github/workflows/CI.yml

@@ -28,7 +28,7 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - "1"
+          - "1.10"
         os:
           - ubuntu-latest
           - macos-latest
@@ -65,7 +65,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        version: ["1"]
+        version: ["1.10"]
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2
@@ -113,5 +113,3 @@ jobs:
 
 env:
   BACKEND_GROUP: "CPU"
-  RETESTITEMS_NWORKERS: 4
-  RETESTITEMS_NWORKER_THREADS: 2

.github/workflows/Documentation.yml

@@ -15,7 +15,7 @@ jobs:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@latest
         with:
-          version: '1'
+          version: '1.10'
       - name: Install dependencies
         run: |
           import Pkg

docs/src/index.md

@@ -14,9 +14,9 @@ Pkg.add("NeuralOperators")
 
 Currently provided operator architectures are :
 
-  - [Fourier Neural Operators (FNOs)](tutorials/fno.md)
-  - [DeepONets](tutorials/deeponet.md)
-  - [Nonlinear Manifold Decoders for Operator Learning (NOMADs)](tutorials/nomad.md)
+  - [Fourier Neural Operators (FNOs)](models/fno.md)
+  - [DeepONets](models/deeponet.md)
+  - [Nonlinear Manifold Decoders for Operator Learning (NOMADs)](models/nomad.md)
 
 ## Reproducibility
 

docs/src/tutorials/burgers.md

@@ -79,19 +79,23 @@ function loss_function(model, ps, st, ((v, y), u))
     return MAELoss()(û, u), stₙ, (;)
 end
 
-begin
-    train_state = Training.TrainState(deeponet, ps, st, Adam(0.0001f0))
-    for epoch in 1:5000
-        grads, loss, stats, train_state = Training.single_train_step!(
-            AutoZygote(), loss_function, ((x_data_dev, grid_dev), y_data_dev), train_state)
+function train_model!(model, ps, st, data; epochs=5000)
+    train_state = Training.TrainState(model, ps, st, Adam(0.0001f0))
+
+    for epoch in 1:epochs
+        _, loss, _, train_state = Training.single_train_step!(
+            AutoZygote(), loss_function, data, train_state)
 
-        if epoch % 25 == 1 || epoch == 400
+        if epoch % 25 == 1 || epoch == epochs
             @printf("Epoch %d: loss = %.6e\n", epoch, loss)
         end
     end
-    ps_trained = train_state.parameters
-    st_trained = train_state.states
+
+    return train_state.parameters, train_state.states
 end
+
+ps_trained, st_trained = train_model!(
+    deeponet, ps, st, ((x_data_dev, grid_dev), y_data_dev))
 ```
 
 ## Plotting

src/NeuralOperators.jl

@@ -11,7 +11,7 @@ using Lux
 using LuxCore: LuxCore, AbstractLuxLayer, AbstractLuxContainerLayer, AbstractLuxWrapperLayer
 using LuxLib: batched_matmul
 using MLDataDevices: AbstractDevice, AbstractGPUDevice
-using NNlib: NNlib, batched_adjoint
+using NNlib: NNlib
 
 const BoolLike = Union{Bool, StaticBool, Val{true}, Val{false}}
 const CRC = ChainRulesCore