From bdda313c8a0325c6ccda45d8999ff16cf46bce17 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Tue, 17 Sep 2024 11:18:47 -0400
Subject: [PATCH] docs: simplify getting started docs

---
 .buildkite/documentation.yml             | 14 +---
 docs/src/components/AsideTrustees.vue    | 18 -----
 docs/src/introduction/index.md           | 88 +++++++++++++++---------
 examples/OptimizationIntegration/main.jl |  2 +-
 4 files changed, 60 insertions(+), 62 deletions(-)

diff --git a/.buildkite/documentation.yml b/.buildkite/documentation.yml
index 2a4398a12..c9253aeb6 100644
--- a/.buildkite/documentation.yml
+++ b/.buildkite/documentation.yml
@@ -8,11 +8,6 @@ steps:
         plugins:
           - JuliaCI/julia#v1:
               version: "1"
-          - JuliaCI/julia-coverage#v1:
-              codecov: true
-              dirs:
-                - src
-                - ext
         command: julia --code-coverage=user --color=yes --project=docs --threads=auto docs/tutorials.jl
         env:
           TUTORIAL_BACKEND_GROUP: "CUDA"
@@ -25,7 +20,7 @@ steps:
           - "docs/src/tutorials/advanced/**/*"
           - "tutorial_deps/*"
           - "**/*.cov"
-        timeout_in_minutes: 60
+        timeout_in_minutes: 120
 
       - label: "Tutorial Build [%N/%t] CPU Runners"
         if: build.message !~ /\[skip docs\]/ && !build.pull_request.draft
@@ -34,11 +29,6 @@ steps:
         plugins:
           - JuliaCI/julia#v1:
               version: "1"
-          # - JuliaCI/julia-coverage#v1:
-          #     codecov: true
-          #     dirs:
-          #       - src
-          #       - ext
         command: julia --code-coverage=user --color=yes --project=docs --threads=auto docs/tutorials.jl
         env:
           TUTORIAL_BACKEND_GROUP: "CPU"
@@ -52,7 +42,7 @@ steps:
           - "docs/src/tutorials/advanced/**/*"
           - "tutorial_deps/*"
           - "**/*.cov"
-        timeout_in_minutes: 60
+        timeout_in_minutes: 120
 
       - label: "Final Documentation Build"
         depends_on:
diff --git a/docs/src/components/AsideTrustees.vue b/docs/src/components/AsideTrustees.vue
index 3b4ac349a..1b13417b5 100644
--- a/docs/src/components/AsideTrustees.vue
+++ b/docs/src/components/AsideTrustees.vue
@@ -10,24 +10,6 @@
             <p class="extra-info">Machine Learning</p>
         </span>
     </a>
-
-    <a class="enjoyer" href="https://juliagni.github.io/GeometricMachineLearning.jl/latest/" target="_blank">
-        <img width="32" height="32" src="https://juliagni.github.io/GeometricMachineLearning.jl/latest/assets/logo-dark.png" />
-        <span>
-            <p class="extra-info">Structure Preserving</p>
-            <p class="heading">GeometricML.jl</p>
-            <p class="extra-info">Machine Learning</p>
-        </span>
-    </a>
-
-    <a class="enjoyer" href="https://una-auxme.github.io/MeshGraphNets.jl/dev/" target="_blank">
-        <img width="32" height="32" src="https://raw.githubusercontent.com/una-auxme/MeshGraphNets.jl/main/logo/meshgraphnetsjl_logo.png" />
-        <span>
-            <p class="extra-info">Physical Systems</p>
-            <p class="heading">MeshGraphNets.jl</p>
-            <p class="extra-info">Graph Neural Nets</p>
-        </span>
-    </a>
 </template>
 
 <style>
diff --git a/docs/src/introduction/index.md b/docs/src/introduction/index.md
index 43a1c8717..b67ce45d6 100644
--- a/docs/src/introduction/index.md
+++ b/docs/src/introduction/index.md
@@ -25,7 +25,8 @@ Pkg.add("Lux")
 
 ```@example quickstart
 using Lux, Random, Optimisers, Zygote
-# using LuxCUDA, AMDGPU, Metal, oneAPI # Optional packages for GPU support
+using LuxCUDA # For CUDA support
+# using AMDGPU, Metal, oneAPI # Other pptional packages for GPU support
 ```
 
 We take randomness very seriously
@@ -43,15 +44,16 @@ Build the model
 model = Chain(Dense(128, 256, tanh), Chain(Dense(256, 1, tanh), Dense(1, 10)))
 ```
 
-Models don't hold parameters and states so initialize them. From there on, we just use our
-standard AD and Optimisers API.
+Models don't hold parameters and states so initialize them. From there on, we can just use
+our standard AD and Optimisers API. However, here we will show how to use Lux's Training
+API that provides an uniform API over all supported AD systems.
 
 ```@example quickstart
 # Get the device determined by Lux
 dev = gpu_device()
 
 # Parameter and State Variables
-ps, st = Lux.setup(rng, model) .|> dev
+ps, st = Lux.setup(rng, model) |> dev
 
 # Dummy Input
 x = rand(rng, Float32, 128, 2) |> dev
@@ -60,21 +62,30 @@ x = rand(rng, Float32, 128, 2) |> dev
 y, st = Lux.apply(model, x, ps, st)
 
 # Gradients
-## Pullback API to capture change in state
-(l, st_), pb = pullback(Lux.apply, model, x, ps, st)
-gs = pb((one.(l), nothing))[3]
+## First construct a TrainState
+train_state = Lux.Training.TrainState(model, ps, st, Adam(0.0001f0))
 
-# Optimization
-st_opt = Optimisers.setup(Adam(0.0001f0), ps)
-st_opt, ps = Optimisers.update(st_opt, ps, gs)  # or Optimisers.update!(st_opt, ps, gs)
+## We can compute the gradients using Training.compute_gradients
+gs, loss, stats, train_state = Lux.Training.compute_gradients(AutoZygote(), MSELoss(),
+    (x, dev(rand(rng, Float32, 10, 2))), train_state)
+
+## Optimization
+train_state = Training.apply_gradients!(train_state, gs) # or Training.apply_gradients (no `!` at the end)
+
+# Both these steps can be combined into a single call
+gs, loss, stats, train_state = Training.single_train_step!(AutoZygote(), MSELoss(),
+    (x, dev(rand(rng, Float32, 10, 2))), train_state)
 ```
 
 ## Defining Custom Layers
 
 ```@example custom_compact
 using Lux, Random, Optimisers, Zygote
-# using LuxCUDA, AMDGPU, Metal, oneAPI # Optional packages for GPU support
-using Printf  # For pretty printing
+using LuxCUDA # For CUDA support
+# using AMDGPU, Metal, oneAPI # Other pptional packages for GPU support
+using Printf # For pretty printing
+
+dev = gpu_device()
 ```
 
 We will define a custom MLP using the `@compact` macro. The macro takes in a list of
@@ -86,9 +97,9 @@ n_in = 1
 n_out = 1
 nlayers = 3
 
-model = @compact(w1=Dense(n_in, 128),
-    w2=[Dense(128, 128) for i in 1:nlayers],
-    w3=Dense(128, n_out),
+model = @compact(w1=Dense(n_in => 32),
+    w2=[Dense(32 => 32) for i in 1:nlayers],
+    w3=Dense(32 => n_out),
     act=relu) do x
     embed = act(w1(x))
     for w in w2
@@ -102,26 +113,41 @@ end
 We can initialize the model and train it with the same code as before!
 
 ```@example custom_compact
-ps, st = Lux.setup(Xoshiro(0), model)
+rng = Random.default_rng()
+Random.seed!(rng, 0)
+
+ps, st = Lux.setup(Xoshiro(0), model) |> dev
 
-model(randn(n_in, 32), ps, st)  # 1×32 Matrix as output.
+x = rand(rng, Float32, n_in, 32) |> dev
 
-x_data = collect(-2.0f0:0.1f0:2.0f0)'
+model(x, ps, st)  # 1×32 Matrix and updated state as output.
+
+x_data = reshape(collect(-2.0f0:0.1f0:2.0f0), 1, :) |> dev
 y_data = 2 .* x_data .- x_data .^ 3
-st_opt = Optimisers.setup(Adam(), ps)
 
-for epoch in 1:1000
-    global st  # Put this in a function in real use-cases
-    (loss, st), pb = Zygote.pullback(ps) do p
-        y, st_ = model(x_data, p, st)
-        return sum(abs2, y .- y_data), st_
+function train_model!(model, ps, st, x_data, y_data)
+    train_state = Lux.Training.TrainState(model, ps, st, Adam(0.001f0))
+
+    for iter in 1:1000
+        _, loss, _, train_state = Lux.Training.single_train_step!(AutoZygote(), MSELoss(),
+            (x_data, y_data), train_state)
+        if iter % 100 == 1 || iter == 1000
+            @printf "Iteration: %04d \t Loss: %10.9g\n" iter loss
+        end
     end
-    gs = only(pb((one(loss), nothing)))
-    epoch % 100 == 1 && @printf "Epoch: %04d \t Loss: %10.9g\n" epoch loss
-    Optimisers.update!(st_opt, ps, gs)
+
+    return model, ps, st
 end
+
+train_model!(model, ps, st, x_data, y_data)
+nothing #hide
 ```
 
+!!! tip "Training with Optimization.jl"
+
+    If you are coming from the SciML ecosystem and want to use Optimization.jl, please
+    refer to the [Optimization.jl Tutorial](@ref Optimization-Lux-Tutorial).
+
 ## Additional Packages
 
 `LuxDL` hosts various packages that provide additional functionality for Lux.jl. All
@@ -133,7 +159,7 @@ You can install all those packages via `import Pkg; Pkg.add(<package name>)`.
 
 GPU Support for Lux.jl requires loading additional packages:
 
-* [`LuxCUDA.jl`](https://github.com/LuxDL/LuxCUDA.jl) for CUDA support.
-* [`AMDGPU.jl`](https://github.com/JuliaGPU/AMDGPU.jl) for AMDGPU support.
-* [`Metal.jl`](https://github.com/JuliaGPU/Metal.jl) for Apple Metal support.
-* [`oneAPI.jl`](https://github.com/JuliaGPU/oneAPI.jl) for oneAPI support.
+- [`LuxCUDA.jl`](https://github.com/LuxDL/LuxCUDA.jl) for CUDA support.
+- [`AMDGPU.jl`](https://github.com/JuliaGPU/AMDGPU.jl) for AMDGPU support.
+- [`Metal.jl`](https://github.com/JuliaGPU/Metal.jl) for Apple Metal support.
+- [`oneAPI.jl`](https://github.com/JuliaGPU/oneAPI.jl) for oneAPI support.
diff --git a/examples/OptimizationIntegration/main.jl b/examples/OptimizationIntegration/main.jl
index b0b9bfeb7..b2ba1c1bc 100644
--- a/examples/OptimizationIntegration/main.jl
+++ b/examples/OptimizationIntegration/main.jl
@@ -1,4 +1,4 @@
-# # Training Lux Models using Optimization.jl
+# # [Training Lux Models using Optimization.jl](@id Optimization-Lux-Tutorial)
 
 # Lux's native [Training.TrainState](@ref) is a great API for gradient-based learning of
 # neural networks, however, it is geared towards using `Optimisers.jl` as the backend.