### Install and Load Flux.jl Packages

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/models/quickstart.md

Installs necessary packages like Flux, CUDA, and ProgressMeter, then loads them. Includes optional CUDA setup for GPU acceleration.

```julia
# Install everything, including CUDA, and load packages:
using Pkg; Pkg.add(["Flux", "CUDA", "cuDNN", "ProgressMeter"])
using Flux, Statistics, ProgressMeter
using CUDA  # optional
device = gpu_device()  # function to move data and model to the GPU
```

--------------------------------

### Install MPI.jl mpiexecjl

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/gpu.md

Install the mpiexecjl executable for MPI.jl. This is a prerequisite for distributed training.

```julia-repl
julia> using MPI

julia> MPI.install_mpiexecjl()
```

--------------------------------

### Setup Flux Code Formatting Tooling

Source: https://github.com/fluxml/flux.jl/blob/master/CONTRIBUTING.md

Run this script to set up the necessary tooling for Flux's code formatting style, including a pre-commit hook.

```julia
julia dev/setup.jl

```

--------------------------------

### List all available tests

Source: https://github.com/fluxml/flux.jl/blob/master/AGENTS.md

Use this command to list all available tests in the Flux.jl project. It requires Julia to be installed and the project to be cloned.

```bash
julia --project=test/ test/runtests.jl --list
```

--------------------------------

### Example of Custom Affine Layer Usage

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/tutorials/custom_layers.md

Demonstrates the default behavior of `Flux.trainable` and the effect of custom `trainable` overload for an `Affine` layer.

```julia
a = Affine(Float32[1 2; 3 4; 5 6], Float32[7, 8, 9])

Flux.trainable(a) # default behavior

Flux.trainable(a::Affine) = (; W = a.W)  # returns a NamedTuple using the field's name

Flux.trainable(a)
```

--------------------------------

### Automating Model Building with outputsize

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/outputsize.md

This example demonstrates how to automate model building by using `Flux.outputsize` to compute the output dimensions of convolutional layers and then defining an appropriate Dense layer.

```julia
using Flux

function make_model(width, height, inchannels = 1, nclasses = 10;
                    layer_config = [16, 16, 32, 64])
  # construct a vector of layers:
  conv_layers = []
  push!(conv_layers, Conv((5, 5), inchannels => layer_config[1], relu, pad=SamePad()))
  for (inch, outch) in zip(layer_config, layer_config[2:end])
    push!(conv_layers, Conv((3, 3), inch => outch, sigmoid, stride=2))
  end

  # compute the output dimensions after these conv layers:
  conv_outsize = Flux.outputsize(conv_layers, (width, height, inchannels); padbatch=true)

  # use this to define appropriate Dense layer:
  last_layer = Dense(prod(conv_outsize) => nclasses)
  return Chain(conv_layers..., Flux.flatten, last_layer)
end

m = make_model(28, 28, 3, layer_config = [9, 17, 33, 65])

Flux.outputsize(m, (28, 28, 3, 42)) == (10, 42) == size(m(randn(Float32, 28, 28, 3, 42)))

# output

true
```

--------------------------------

### GPU Gradient Calculation and Optimizer Setup

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/gpu.md

Illustrates how Flux's gradient calculation and training functions work seamlessly with GPU models and data. It's recommended to move the model to the GPU before setting up the optimizer.

```julia
grads = Flux.gradient((f,x) -> sum(abs2, f(x)), model, x)  # on CPU
c_grads = Flux.gradient((f,x) -> sum(abs2, f(x)), c_model, cx)  # same result, all on GPU

c_opt = Flux.setup(Adam(), c_model)  # setup optimiser after moving model to GPU

Flux.update!(c_opt, c_model, c_grads[1])  # mutates c_model but not model
```

--------------------------------

### Get Device

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/data/mldatadevices.md

Retrieves the current active device.

```APIDOC
## MLDataDevices.get_device

### Description
Returns the currently active computational device.

### Method
`get_device()`

### Endpoint
N/A (Function call)

### Parameters
None

### Request Example
```julia
using MLDataDevices
current_device = get_device()
```

### Response
#### Success Response
- `Device`: An object representing the current active device.
```

--------------------------------

### Dense Layer Example

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/models/layers.md

The Dense layer is a fundamental building block for neural networks, featuring an activation function, an initialization function for weights, and an optional bias term. It is annotated with @layer for automatic setup and GPU movement.

```julia
Dense(in => out, activation; init = Flux.glorot_uniform, bias = Flux.zeros32, init_bias = Flux.zeros32)
```

--------------------------------

### Checkpoint Model Training with JLD2

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/saving.md

Periodically saves the model's state during training using JLD2.jl to allow resumption if training is interrupted. This example updates a single checkpoint file.

```jldoctest saving
julia> using Flux: throttle

julia> using JLD2

julia> m = Chain(Dense(10 => 5, relu), Dense(5 => 2))
Chain(
  Dense(10 => 5, relu),
  Dense(5 => 2),
)                   # Total: 4 arrays, 67 parameters, 476 bytes.

julia> for epoch in 1:10
          # ... train model ...
          jldsave("model-checkpoint.jld2", model_state = Flux.state(m))
       end;
```

--------------------------------

### Build a Model using Flux's Chain Layer

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/models/basics.md

Demonstrates how to construct a neural network model using Flux's pre-defined `Chain` layer, which sequentially applies a list of layers. This example uses `Dense` layers and the `only` function.

```jldoctest
model3 = Chain(Dense(1 => 20, σ), Dense(20 => 1), only)
```

--------------------------------

### Get Device Type

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/data/mldatadevices.md

Determines the type of a given device.

```APIDOC
## MLDataDevices.get_device_type

### Description
Returns the type of the given device (e.g., CPU, GPU).

### Method
`get_device_type(device)`

### Endpoint
N/A (Function call)

### Parameters
- **device** (Device) - Required - The device to query.

### Request Example
```julia
using MLDataDevices
device_type = get_device_type(cpu_device())
```

### Response
#### Success Response
- `Symbol`: A symbol representing the device type (e.g., `:cpu`, `:gpu`).
```

--------------------------------

### Nested Structural Gradients

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/models/basics.md

Demonstrates calculating structural gradients for deeply nested structures, such as nested tuples and structs. This example shows how to access gradients for inner parameters like `poly4.inner.θ3`.

```jldoctest poly
julia> grad4 = gradient(|>, 5, poly4)
(1.0, (outer = (θ3 = [1.0, 17.5, 306.25],), inner = (θ3 = [0.5, 2.5, 12.5],)))
```

--------------------------------

### Manually Format Flux Codebase

Source: https://github.com/fluxml/flux.jl/blob/master/CONTRIBUTING.md

Execute this command to manually format the entire Flux codebase according to the project's style guide.

```julia
julia dev/flux_format.jl --verbose .

```

--------------------------------

### Run tests with CUDA backend enabled

Source: https://github.com/fluxml/flux.jl/blob/master/AGENTS.md

Execute tests using the CUDA backend by setting the FLUX_TEST_CUDA environment variable to true. Ensure CUDA is properly installed and configured.

```bash
FLUX_TEST_CUDA=true julia --project=test/ test/runtests.jl
```

--------------------------------

### Chain Layer Example

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/models/layers.md

Chain combines layers sequentially, applying them in the order they are provided, equivalent to function composition. It does not contain trainable parameters itself but orchestrates the flow through its constituent layers.

```julia
Chain(layer1, layer2, layer3)
```

--------------------------------

### Initialise Momentum Optimizer

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/training/training.md

Initialises the optimiser state for a given model using the Momentum optimiser. This is typically done before starting the training loop.

```julia
opt_state = Flux.setup(Momentum(0.01, 0.9), model)
```

--------------------------------

### Move Flux Model and Data to GPU

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/gpu.md

Shows how to move a Flux model and its associated data to the GPU using the `cu` function. Ensure CUDA.jl and cuDNN are installed and loaded.

```julia
using Pkg; Pkg.add(["CUDA", "cuDNN"])  # do this once

using Flux, CUDA
CUDA.allowscalar(false)  # recommended

model = Dense(W, true, tanh)  # wrap the same matrix W in a Flux layer
model(x) ≈ y                  # same result, still on CPU

c_model = cu(model)  # move all the arrays within model to the GPU
c_model(cx)          # computation on the GPU
```

--------------------------------

### Optimisers.setup

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/training/reference.md

Sets up the optimizer state for a given model. This is the general function from the Optimisers.jl package.

```APIDOC
## Optimisers.setup

### Description
Sets up the optimizer state for a given model. This is the general function from the Optimisers.jl package, which Flux.Train.setup wraps.

### Method
```
Optimisers.setup(opt, model)
```

### Parameters
* `opt`: The optimizer rule (e.g., `Adam()`).
* `model`: The model for which to set up the optimizer state.

### Returns
The initial state of the optimizer.
```

--------------------------------

### Initialize Training Components

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/models/overview.md

Sets up the optimiser and training data before calling Flux.train!.

```julia
using Flux: train!

opt = Descent()
data = [(x_train, y_train)]
```

--------------------------------

### Create a Simple Dataset

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/training/training.md

Demonstrates how to create a basic dataset as a vector of tuples, where each tuple contains an input and its corresponding label.

```julia
x = randn(28, 28)
y = rand(10)
data = [(x, y)]
```

--------------------------------

### Default Device RNG

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/data/mldatadevices.md

Gets the default random number generator for the current device.

```APIDOC
## MLDataDevices.default_device_rng

### Description
Retrieves the default random number generator (RNG) associated with the current active device.

### Method
`default_device_rng()`

### Endpoint
N/A (Function call)

### Parameters
None

### Request Example
```julia
using MLDataDevices
rng = default_device_rng()
```

### Response
#### Success Response
- `AbstractRNG`: An object representing the default random number generator for the device.
```

--------------------------------

### Set up and Synchronize Distributed Optimizer

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/gpu.md

Create a DistributedOptimizer with your chosen optimization algorithm (e.g., Adam) and synchronize it across all processes. This ensures consistent optimization state.

```julia-repl
julia> using Optimisers

julia> opt = DistributedUtils.DistributedOptimizer(backend, Optimisers.Adam(0.001f0))
DistributedOptimizer{MPIBackend{Comm}}(MPIBackend{Comm}(Comm(1140850688)), Adam(0.001, (0.9, 0.999), 1.0e-8))

julia> st_opt = Optimisers.setup(opt, model)
(layers = ((weight = Leaf(DistributedOptimizer{MPIBackend{Comm}}(MPIBackend{Comm}(Comm(1140850688)), Adam(0.001, (0.9, 0.999), 1.0e-8)), (Float32[0.0; 0.0; … ; 0.0; 0.0;;], Float32[0.0; 0.0; … ; 0.0; 0.0;;], (0.9, 0.999))), bias = Leaf(DistributedOptimizer{MPIBackend{Comm}}(MPIBackend{Comm}(Comm(1140850688)), Adam(0.001, (0.9, 0.999), 1.0e-8)), (Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0  …  0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0  …  0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], (0.9, 0.999))), σ = ()), (weight = Leaf(DistributedOptimizer{MPIBackend{Comm}}(MPIBackend{Comm}(Comm(1140850688)), Adam(0.001, (0.9, 0.999), 1.0e-8)), (Float32[0.0 0.0 … 0.0 0.0], Float32[0.0 0.0 … 0.0 0.0], (0.9, 0.999))), bias = Leaf(DistributedOptimizer{MPIBackend{Comm}}(MPIBackend{Comm}(Comm(1140850688)), Adam(0.001, (0.9, 0.999), 1.0e-8)), (Float32[0.0], Float32[0.0], (0.9, 0.999))), σ = ())),)

julia> st_opt = DistributedUtils.synchronize!!(backend, st_opt; root=0)
```

--------------------------------

### focal_loss

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/models/losses.md

Focal loss function. An extension of cross-entropy that down-weights easy examples and focuses on hard negatives.

```APIDOC
## focal_loss

### Description
Focal loss function. An extension of cross-entropy that down-weights easy examples and focuses on hard negatives.

### Method
`loss(ŷ, y, α, γ, agg)`

### Parameters
- **ŷ** (Any) - Predicted probabilities.
- **y** (Any) - Target values.
- **α** (Real) - Weighting factor for the positive class.
- **γ** (Real) - Focusing parameter.
- **agg** (Any) - Optional aggregation function (defaults to `mean`).

### Response
- Returns the calculated focal loss.
```

--------------------------------

### binary_focal_loss

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/models/losses.md

Binary focal loss function. An extension of binary cross-entropy that down-weights easy examples and focuses on hard negatives.

```APIDOC
## binary_focal_loss

### Description
Binary focal loss function. An extension of binary cross-entropy that down-weights easy examples and focuses on hard negatives.

### Method
`loss(ŷ, y, α, γ, agg)`

### Parameters
- **ŷ** (Any) - Predicted probabilities.
- **y** (Any) - Target binary labels (0 or 1).
- **α** (Real) - Weighting factor for the positive class.
- **γ** (Real) - Focusing parameter.
- **agg** (Any) - Optional aggregation function (defaults to `mean`).

### Response
- Returns the calculated binary focal loss.
```

--------------------------------

### Manual Training Loop

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/training/training.md

This snippet shows the fundamental steps of a training loop: initializing the optimizer state, iterating through training data, calculating gradients, and updating model parameters.

```julia
# Initialise the optimiser for this model:
opt_state = Flux.setup(rule, model)

for data in train_set
  # Unpack this element (for supervised training):
  input, label = data

  # Calculate the gradient of the objective
  # with respect to the parameters within the model:
  grads = Flux.gradient(model) do m
      result = m(input)
      loss(result, label)
  end

  # Update the parameters so as to reduce the objective,
  # according the chosen optimisation rule:
  Flux.update!(opt_state, model, grads[1])
end
```

--------------------------------

### Flux.Train.setup

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/training/reference.md

Sets up the optimizer state for a given model. This function is specific to Flux and ensures compatibility with its parameter structures.

```APIDOC
## Flux.Train.setup

### Description
Sets up the optimizer state for a given model, ensuring compatibility with Flux's parameter structures. This is a Flux-specific wrapper around `Optimisers.setup`.

### Method
```
Flux.Train.setup(opt, model)
```

### Parameters
* `opt`: The optimizer rule (e.g., `Adam()`).
* `model`: The Flux model to set up for optimization.

### Returns
The initial state of the optimizer.
```

--------------------------------

### Test MPI CUDA Awareness

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/gpu.md

Verify if your MPI installation is CUDA-aware. This is a prerequisite for running distributed training on GPUs with CUDA-aware MPI support.

```julia-repl
julia> import Pkg
julia> Pkg.test("MPI"; test_args=["--backend=CUDA"])
```

--------------------------------

### Customize Trainable Parameters

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/tutorials/custom_layers.md

Customize which fields of a layer are considered trainable by overloading the `trainable` function. This example makes only the `W` field trainable.

```julia
Flux.trainable(a::Affine) = (; W = a.W)
```

--------------------------------

### Configure Initialisation Function with Gain

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/utilities.md

Creates a `Dense` layer with weights initialised using `glorot_uniform`, configured with a `gain` of 2.

```jldoctest
julia> Dense(4 => 5, tanh; init=Flux.glorot_uniform(gain=2))
Dense(4 => 5, tanh)  # 25 parameters
```

--------------------------------

### Initialize Distributed Backend

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/gpu.md

Initialize the distributed backend, preferably NCCLBackend if NCCL.jl is loaded, for GPU-accelerated distributed training. Ensure scalar operations are disallowed on CUDA devices.

```julia-repl
julia> using Flux, MPI, NCCL, CUDA

julia> CUDA.allowscalar(false)

julia> DistributedUtils.initialize(NCCLBackend)

julia> backend = DistributedUtils.get_distributed_backend(NCCLBackend)
NCCLBackend{Communicator, MPIBackend{MPI.Comm}}(Communicator(Ptr{NCCL.LibNCCL.ncclComm} @0x000000000607a660), MPIBackend{MPI.Comm}(MPI.Comm(1140850688)))
```

--------------------------------

### Prepare Data Loader and Optimizer

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/models/quickstart.md

Prepares the data for training by one-hot encoding the targets and creating a DataLoader for batching. Sets up the Adam optimizer with a specified learning rate.

```julia
# The model encapsulates parameters, randomly initialised. Its initial output is:
out1 = model(noisy |> device)    # 2×1000 Matrix{Float32}, or CuArray{Float32}
probs1 = softmax(out1) |> cpu    # normalise to get probabilities (and move off GPU)

# To train the model, we use batches of 64 samples, and one-hot encoding:
target = Flux.onehotbatch(truth, [true, false])                   # 2×1000 OneHotMatrix
loader = Flux.DataLoader((noisy, target), batchsize=64, shuffle=true);

opt_state = Flux.setup(Flux.Adam(0.01), model)  # will store optimiser momentum, etc.
```

--------------------------------

### Format code in src directory

Source: https://github.com/fluxml/flux.jl/blob/master/AGENTS.md

Format all Julia code within the 'src' directory using JuliaFormatter. This command requires the JuliaFormatter package to be installed.

```bash
julia -e 'using JuliaFormatter; format("src")'
```

--------------------------------

### Run all CPU tests with quick-fail enabled

Source: https://github.com/fluxml/flux.jl/blob/master/AGENTS.md

Run all CPU tests and stop execution immediately upon the first test failure. This is useful for rapid debugging.

```bash
julia --project=test/ test/runtests.jl --quickfail
```

--------------------------------

### Train with Duplicated Model and Enzyme.jl

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/training/reference.md

Example of using `Flux.train!` with a duplicated model and Enzyme.jl for automatic differentiation. Requires `dup_model` to be set up with `Duplicated(model)`.

```julia
julia> opt_state = Flux.setup(Adam(0), model);

julia> Flux.train!((m,x,y) -> sum(abs2, m(x) .- y), dup_model, [(x1, y1)], opt_state)
```

--------------------------------

### Supported GPU Backends

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/data/mldatadevices.md

Lists the GPU backends supported by the system.

```APIDOC
## MLDataDevices.supported_gpu_backends

### Description
Returns a list of GPU backends that are supported and available on the current system.

### Method
`supported_gpu_backends()`

### Endpoint
N/A (Function call)

### Parameters
None

### Request Example
```julia
using MLDataDevices
backends = supported_gpu_backends()
```

### Response
#### Success Response
- `Vector{Symbol}`: A vector of symbols representing the supported GPU backends (e.g., `[:cuda, :rocm]`).
```

--------------------------------

### Define a Polynomial Function with Global Parameters

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/models/basics.md

Defines a simple polynomial function where parameters are stored in a global variable. This is a basic example to illustrate parameter dependence.

```jldoctest
θ = [10, 1, 0.1]

poly1(x::Real) = θ[1] + θ[2]*x + θ[3]*x^2

poly1(5) == 17.5  # true

# output

true
```

--------------------------------

### Enable Flux CUDA-Aware MPI Preference

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/gpu.md

Set a local preference to enable Flux's support for CUDA-aware MPI. This should be done after confirming your MPI installation is CUDA-aware.

```julia-repl
julia> using Preferences
julia> set_preferences!("Flux", "FluxDistributedMPICUDAAware" => true)
```

--------------------------------

### Flux.@layer Macro

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/models/functors.md

The Flux.@layer macro is recommended for defining layers. It allows Flux.setup to access parameters within the layer and enables functions like gpu to move them to the GPU. It also overloads printing and offers a way to define `trainable`.

```APIDOC
## Flux.@layer

### Description
This macro is used to define layers in Flux. It enables parameter access for functions like `Flux.setup` and `gpu`, and provides enhanced printing capabilities.

### Usage
```julia
Flux.@layer MyLayer
```
```

--------------------------------

### Is Loaded

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/data/mldatadevices.md

Checks if a specific GPU backend is loaded and available.

```APIDOC
## MLDataDevices.loaded

### Description
Checks if a specific GPU backend (e.g., CUDA) is loaded and available for use.

### Method
`loaded(backend::Symbol)`

### Endpoint
N/A (Function call)

### Parameters
- **backend** (Symbol) - Required - The GPU backend to check (e.g., `:cuda`).

### Request Example
```julia
using MLDataDevices
cuda_loaded = loaded(:cuda)
```

### Response
#### Success Response
- `Bool`: `true` if the backend is loaded, `false` otherwise.
```

--------------------------------

### Gradient Clipping

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/training/optimisers.md

Provides gradient clipping methods, useful for mitigating the exploding gradient problem, particularly in recurrent neural networks. Examples show integration with other optimisers.

```APIDOC
## Gradient Clipping

Gradient clipping is useful for training recurrent neural networks, which have a tendency to suffer from the exploding gradient problem. An example usage is

```julia
opt = OptimiserChain(ClipGrad(1e-3), Adam(1e-3))
```

```@docs
Optimisers.ClipGrad
Optimisers.ClipNorm
```
```

--------------------------------

### Select and Inspect a GPU Device

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/gpu.md

Selects a specific GPU device by its ID (starting from 1 for gpu_device!) and displays its properties. Note the difference in indexing between CUDA.devices() and gpu_device!.

```julia-repl
julia> device0 = gpu_device(1)
(::CUDADevice{CuDevice}) (generic function with 4 methods)

julia> device0.device
CuDevice(0): NVIDIA TITAN RTX
```

--------------------------------

### Check Metal GPU Functionality

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/gpu.md

Verify if the Metal GPU is functional. Ensure the Metal package is imported.

```julia-repl
julia> using Metal

julia> Metal.functional()
true
```

--------------------------------

### Inspect Model Parameters

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/models/overview.md

Shows the initial weights and biases of a dense layer model before training.

```julia
predict.weight
1×1 Matrix{Float32}:
 0.9066542
```

```julia
predict.bias
1-element Vector{Float32}:
 0.0
```

--------------------------------

### Composing Optimisers with OptimiserChain

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/training/optimisers.md

Demonstrates how to compose optimisers using `OptimiserChain` for sequential application of updates, such as adding weight decay. The composed optimiser can be used like any standard optimiser.

```APIDOC
## Composing Optimisers

Flux (through Optimisers.jl) defines a special kind of optimiser called `OptimiserChain` which takes in arbitrary optimisers as input. Its behaviour is similar to the usual optimisers, but differs in that it acts by calling the optimisers listed in it sequentially. Each optimiser produces a modified gradient that will be fed into the next, and the resultant update will be applied to the parameter as usual. A classic use case is where adding decays is desirable. Optimisers.jl defines the basic decay corresponding to an $L_2$ regularization in the loss as `WeightDecay`.

```julia
opt = OptimiserChain(WeightDecay(1e-4), Descent())
```

Here we apply the weight decay to the `Descent` optimiser. 
The resulting optimiser `opt` can be used as any optimiser.

```julia
w = [randn(10, 10), randn(10, 10)]
opt_state = Flux.setup(opt, w)

loss(w, x) = Flux.mse(w[1] * x, w[2] * x)

loss(w, rand(10)) # around 0.9

for t = 1:10^5
  g = gradient(w -> loss(w[1], w[2], rand(10)), w)
  Flux.update!(opt_state, w, g)
end

loss(w, rand(10)) # around 0.9
```

It is possible to compose optimisers for some added flexibility.

```@docs
Optimisers.OptimiserChain
```
```

--------------------------------

### Load and Prepare Boston Housing Dataset

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/tutorials/linear_regression.md

Loads the Boston Housing dataset, converts it to Float32, and splits it into training and testing sets.

```jldoctest
julia> dataset = BostonHousing();

julia> x, y = BostonHousing(as_df=false)[:];

julia> x, y = Float32.(x), Float32.(y);
```

```jldoctest
julia> x_train, x_test, y_train, y_test = x[:, 1:400], x[:, 401:end], y[:, 1:400], y[:, 401:end];

julia> x_train |> size, x_test |> size, y_train |> size, y_test |> size
((13, 400), (13, 106), (1, 400), (1, 106))
```

--------------------------------

### Define a Custom Model Structure

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/tutorials/custom_layers.md

Define a custom model by composing existing layers. This example adds the input to the output of a neural network. Use type parameters for efficiency.

```julia
struct CustomModel{T <: Chain}
  chain::T
end

function (m::CustomModel)(x)
  # Arbitrary code can go here, but note that everything will be differentiated.
  # Zygote does not allow some operations, like mutating arrays.

  return m.chain(x) + x
end

# This is optional but recommended for pretty printing and other niceties
Flux.@layer CustomModel
```

--------------------------------

### Basic Flux.jl Model Training and Plotting

Source: https://github.com/fluxml/flux.jl/blob/master/README.md

This snippet demonstrates how to define a simple model, train it using Flux.jl with the Adam optimizer, and then plot the learned function against the true function. It requires the Plots package for visualization.

```julia
using Flux
data = [(x, 2x-x^3) for x in -2:0.1f0:2]

model = let
  w, b, v = (randn(Float32, 23) for _ in 1:3)  # parameters
  x -> sum(v .* tanh.(w*x .+ b))               # callable
end
# model = Chain(vcat, Dense(1 => 23, tanh), Dense(23 => 1, bias=false), only)

opt_state = Flux.setup(Adam(), model)
for epoch in 1:100
  Flux.train!((m,x,y) -> (m(x) - y)^2, model, data, opt_state)
end

using Plots
plot(x -> 2x-x^3, -2, 2, label="truth")
scatter!(model, -2:0.1f0:2, label="learned")
```

--------------------------------

### Flux.sparse_init

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/utilities.md

Initialises weights with a sparse matrix, where only a specified number of weights are non-zero.

```APIDOC
## Function: Flux.sparse_init

### Description
Initialises weights with a sparse matrix, setting a specified number of weights to non-zero values.

### Parameters
- `num_nonzeros` (Integer): The number of non-zero elements to set in the weight matrix.

### Returns
A function that takes dimensions and returns a sparse weight array.
```

--------------------------------

### Use a Custom Model

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/tutorials/custom_layers.md

Instantiate and use a custom model like any other Flux layer. This demonstrates creating a `CustomModel` with a `Chain` and applying it to sample data.

```julia
chain = Chain(Dense(10 => 10, relu), Dense(10 => 10))
model = CustomModel(chain)
model(rand(Float32, 10))
```

--------------------------------

### Manual RNN Cell Implementation

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/models/recurrence.md

Defines a basic RNN cell manually using matrix multiplications and a tanh activation. This serves as a foundational example before introducing Flux's built-in cells.

```julia
output_size = 5
input_size = 2
Wxh = randn(Float32, output_size, input_size)
Whh = randn(Float32, output_size, output_size)
b = zeros(Float32, output_size)

function rnn_cell(x, h)
    h = tanh.(Wxh * x .+ Whh * h .+ b)
    return h, h
end

seq_len = 3
# dummy input data
x = [rand(Float32, input_size) for i = 1:seq_len] 
# random initial hidden state
h0 = zeros(Float32, output_size)

y = []
ht = h0
for xt in x
    yt, ht = rnn_cell(xt, ht)
    y = [y; [yt]]  # concatenate in non-mutating (AD friendly) way
end
```

--------------------------------

### Using Enzyme.jl with Flux Models

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/training/gradients.md

Demonstrates how to wrap a Flux model with Enzyme.Duplicated to allocate space for gradients and then use Flux.gradient to compute them. The gradient is stored within the duplicated model.

```julia-repl
julia> using Flux, Enzyme

julia> model = Chain(Dense(28^2 => 32, sigmoid), Dense(32 => 10), softmax);  # from model zoo

julia> dup_model = Enzyme.Duplicated(model)  # this allocates space for the gradient
Duplicated(
  Chain(
    Dense(784 => 32, σ),                # 25_120 parameters
    Dense(32 => 10),                    # 330 parameters
    NNlib.softmax,
  ),
  # norm(∇) ≈ 0.0f0
)                   # Total: 4 arrays, 25_450 parameters, 199.391 KiB.

julia> x1 = randn32(28*28, 1);  # fake image

julia> y1 = [i==3 for i in 0:9];  # fake label

julia> grads_f = Flux.gradient((m,x,y) -> sum(abs2, m(x) .- y), dup_model, Const(x1), Const(y1))  # uses Enzyme
((layers = ((weight = Float32[-0.010354728 0.032972857 …
    -0.0014538406], σ = nothing), nothing),), nothing, nothing)
```

--------------------------------

### Structural Gradients for Parameterized Functions

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/models/basics.md

Calculates structural gradients for functions with parameters, distinguishing between gradients with respect to the input `x` and gradients with respect to parameters `θ`. This example shows gradients for `poly2` and `poly3`.

```jldoctest poly
julia> grad2 = gradient(poly2, 5, θ)
(2.0, [1.0, 5.0, 25.0])

julia> grad3 = gradient((x,p) -> p(x), 5, poly3s)
(2.0, (θ3 = [1.0, 5.0, 25.0],))
```

--------------------------------

### Custom Training Loop with Zygote.withgradient

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/training/training.md

A flexible custom training loop that uses Zygote.withgradient to compute gradients and the loss value simultaneously. It allows for detailed control over logging, conditional updates, and early stopping.

```julia
opt_state = Flux.setup(Adam(), model)

my_log = []
for epoch in 1:100
  losses = Float32[]
  for (i, data) in enumerate(train_set)
    input, label = data

    val, grads = Flux.withgradient(model) do m
      result = m(input)
      my_loss(result, label)
    end

    push!(losses, val)

    if !isfinite(val)
      @warn "loss is $val on item $i" epoch
      continue
    end

    Flux.update!(opt_state, model, grads[1])
  end

  acc = my_accuracy(model, train_set)
  push!(my_log, (; acc, losses))

  if  acc > 0.95
    println("stopping after $epoch epochs")
    break
  end
end
```

--------------------------------

### Early Stopping with Decreasing Loss

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/training/callbacks.md

Demonstrates using `early_stopping` to halt training when a pseudo-loss function starts increasing for two consecutive steps. The training loop breaks when the `es()` trigger returns true.

```julia
# create a pseudo-loss that decreases for 4 calls, then starts increasing
# we call this like loss()
loss = let t = 0
  () -> begin
    t += 1
    (t - 4) ^ 2
  end
end

# create an early stopping trigger
# returns true when the loss increases for two consecutive steps
es = early_stopping(loss, 2; init_score = 9)

# this will stop at the 6th (4 decreasing + 2 increasing calls) epoch
for epoch in 1:10
  es() && break
end
```

--------------------------------

### Run all CPU tests with verbose output and 4 workers

Source: https://github.com/fluxml/flux.jl/blob/master/AGENTS.md

Execute all CPU tests with verbose output and parallel processing using 4 workers. Ensure Julia is set up correctly.

```bash
julia --project=test/ test/runtests.jl --verbose --jobs=4
```

--------------------------------

### Initialize Model Parameters (Custom)

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/tutorials/linear_regression.md

Initializes the weight matrix `W` from a uniform distribution and the bias `b` to zero for the custom linear regression model. Note that the filter `r"[+-]?([0-9]*[.])?[0-9]+(f[+-]*[0-9])?"` is used to normalize floating-point numbers in the output.

```jldoctest
julia> W = rand(Float32, 1, 1)
1×1 Matrix{Float32}:
 0.99285793
```

```jldoctest
julia> b = [0.0f0]
1-element Vector{Float32}:
 0.0
```

--------------------------------

### GPU Backend

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/data/mldatadevices.md

Sets the backend for GPU operations.

```APIDOC
## MLDataDevices.gpu_backend!

### Description
Sets the backend for GPU operations. This function is typically used to select between different GPU libraries (e.g., CUDA.jl, ROCm.jl).

### Method
`gpu_backend!(backend::Symbol)`

### Endpoint
N/A (Function call)

### Parameters
- **backend** (Symbol) - Required - The GPU backend to use (e.g., `:cuda`, `:rocm`).

### Request Example
```julia
using MLDataDevices
gpu_backend!(:cuda)
```

### Response
#### Success Response
- `Nothing`: Indicates the backend was successfully set.
```

--------------------------------

### Reimplement Join using Flux.Parallel

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/tutorials/custom_layers.md

Show how the custom `Join` layer functionality can be achieved using Flux's built-in `Parallel` layer, demonstrating `Join` as syntactic sugar.

```julia
Join(combine, paths) = Parallel(combine, paths)
Join(combine, paths...) = Join(combine, paths)

# use vararg/tuple version of Parallel forward pass
model = Chain(
              Join(vcat,
                   Chain(Dense(1 => 5, relu), Dense(5 => 1)),
                   Dense(1 => 2),
                   Dense(1 => 1)
                  ),
              Dense(4 => 1)
             ) |> gpu

xs = map(gpu, (rand(1), rand(1), rand(1)))

model(xs)
```

--------------------------------

### Gradient Calculation with Enzyme (Explicit Backend)

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/models/basics.md

Shows how to use `Flux.withgradient` with the Enzyme AD backend explicitly specified. This requires loading the Enzyme package first.

```julia-repl
julia> using Enzyme

julia> Flux.withgradient((x,p) -> p(x), AutoEnzyme(), 5.0, poly3s)
(val = 17.5, grad = (2.0, Poly3{Vector{Float64}}([1.0, 5.0, 25.0])))
```

--------------------------------

### Cosine Annealing Learning Rate Schedule

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/training/training.md

Illustrates how to use a cosine annealing schedule with a Momentum optimiser for dynamic learning rate adjustment during training. Requires importing ParameterSchedulers.jl.

```julia
using ParameterSchedulers

opt_state = Flux.setup(Momentum(), model)
schedule = Cos(λ0 = 1e-4, λ1 = 1e-2, period = 10)
for (eta, epoch) in zip(schedule, 1:100)
  Flux.adjust!(opt_state, eta)
  # your training code here
end
```

--------------------------------

### Loss Function Aggregation Options

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/models/losses.md

Demonstrates various aggregation methods for loss functions using the optional `agg` argument.

```julia
loss(ŷ, y)                         # defaults to `mean`
loss(ŷ, y, agg=sum)                # use `sum` for reduction
loss(ŷ, y, agg=x->sum(x, dims=2))  # partial reduction
loss(ŷ, y, agg=x->mean(w .* x))    # weighted mean
loss(ŷ, y, agg=identity)           # no aggregation.
```

--------------------------------

### Gradient Clipping with OptimiserChain

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/training/training.md

Demonstrates using OptimiserChain for gradient clipping, a regularization technique that limits the magnitude of gradients to prevent exploding gradients. ClipGrad and ClipNorm are available options.

```julia
ClipGrad(0.1)
```

```julia
ClipNorm(1.0)
```

--------------------------------

### Import Flux and Define Target Function

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/models/overview.md

Imports the Flux library and defines the target function `actual(x) = 4x + 2` that the model will approximate. This is the first step in setting up a Flux program.

```jldoctest overview
julia> using Flux

julia> actual(x) = 4x + 2
actual (generic function with 1 method)
```

--------------------------------

### Flux.identity_init

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/utilities.md

Initialises weights with an identity matrix.

```APIDOC
## Function: Flux.identity_init

### Description
Initialises weights with an identity matrix. This is useful for specific layer initialisations where an identity mapping is desired.

### Returns
A function that takes dimensions and returns an identity weight matrix.
```

--------------------------------

### Gradient Calculation with Mooncake

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/models/basics.md

Demonstrates using `Flux.withgradient` with a specified AD backend, Mooncake. Ensure Mooncake is loaded before use.

```jldoctest
julia> using Mooncake

julia> Flux.withgradient((x,p) -> p(x), AutoMooncake(), 5.0, poly3s)
(val = 17.5, grad = (2.0, Poly3{Vector{Float64}}([1.0, 5.0, 25.0])))
```

--------------------------------

### Configure Initialisation Function with RNG

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/utilities.md

Creates a `Dense` layer with weights initialised using `randn32`, using a specific `MersenneTwister` random number generator.

```jldoctest
julia> Dense(4 => 5, tanh; init=Flux.randn32(MersenneTwister(1)))
Dense(4 => 5, tanh)  # 25 parameters
```

--------------------------------

### Prepare Training and Test Data

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/models/overview.md

Generates training and test datasets using the `actual` function. `x_train` and `x_test` are created using `hcat` for input features, and `y_train` and `y_test` are the corresponding target values.

```jldoctest overview
julia> x_train, x_test = hcat(0:5...), hcat(6:10...)
([0 1 … 4 5], [6 7 … 9 10])

julia> y_train, y_test = actual.(x_train), actual.(x_test)
([2 6 … 18 22], [26 30 … 38 42])
```

--------------------------------

### Stateful Cosine Annealing Learning Rate Schedule

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/training/training.md

Demonstrates using a stateful cosine annealing schedule for learning rate adjustment. This approach uses ParameterSchedulers.Stateful and its next! function.

```julia
using ParameterSchedulers: Stateful, next!

schedule = Stateful(Cos(λ0 = 1e-4, λ1 = 1e-2, period = 10))
for epoch in 1:100
  Flux.adjust!(opt_state, next!(schedule))
  # your training code here
end
```

--------------------------------

### Using the train! Function

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/training/training.md

A concise way to perform training using the `train!` function, which abstracts the manual loop. It requires the model, training set, optimizer state, and a loss function defined within a do block.

```julia
train!(model, train_set, opt_state) do m, x, y
  loss(m(x), y)
end
```

--------------------------------

### Initialize Model Parameters

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/tutorials/logistic_regression.md

Initializes the weight matrix (W) and bias vector (b) for a logistic regression model with 4 inputs and 3 outputs. Random values are used for weights, and biases are set to zero.

```julia
julia> W = rand(Float32, 3, 4);

julia> b = [0.0f0, 0.0f0, 0.0f0];
```

--------------------------------

### Zygote.jl Gradient and Derivative Functions

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/training/gradients.md

Core functions for calculating gradients, Jacobians, and Hessians using Zygote.jl.

```julia
Zygote.jacobian(f, args...)
Zygote.withjacobian(f, args...)
Zygote.hessian
Zygote.hessian_reverse
Zygote.diaghessian
Zygote.pullback
```

--------------------------------

### Move Model and Data to GPU

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/guide/gpu.md

Define your neural network model and input data, then move them to the GPU device for accelerated computation.

```julia-repl
julia> model = Chain(Dense(1 => 256, tanh), Dense(256 => 1)) |> gpu
Chain(
  Dense(1 => 256, tanh),                # 512 parameters
  Dense(256 => 1),                      # 257 parameters
)                   # Total: 4 arrays, 769 parameters, 744 bytes.

julia> x = rand(Float32, 1, 16) |> gpu
1×16 CUDA.CuArray{Float32, 2, CUDA.DeviceMemory}:
 0.239324  0.331029  0.924996  0.55593  0.853093  0.874513  0.810269  0.935858  0.477176  0.564591  0.678907  0.729682  0.96809  0.115833  0.66191  0.75822

julia> y = x .^ 3
1×16 CUDA.CuArray{Float32, 2, CUDA.DeviceMemory}:
 0.0137076  0.0362744  0.791443  0.171815  0.620854  0.668804  0.53197  0.819654  0.108651  0.179971  0.312918  0.388508  0.907292  0.00155418  0.29  0.435899
```

--------------------------------

### Standard Optimisers

Source: https://github.com/fluxml/flux.jl/blob/master/docs/src/reference/training/optimisers.md

This section lists the standard optimisers available. Each optimiser can be used with `train!` and other training functions. Refer to the Optimisers.jl documentation for full interface details.

```APIDOC
## Standard Optimisers

Any optimization rule from Optimisers.jl can be used with [`train!`](@ref Flux.Train.train!) and other training functions.

For full details of how the interface works, see the [Optimisers.jl documentation](https://fluxml.ai/Optimisers.jl/).

```@docs
Optimisers.Descent
Optimisers.Momentum
Optimisers.Nesterov
Optimisers.RMSProp
Optimisers.Adam
Optimisers.RAdam
Optimisers.AdaMax
Optimisers.AdaGrad
Optimisers.AdaDelta
Optimisers.AMSGrad
Optimisers.NAdam
Optimisers.AdamW
Optimisers.OAdam
Optimisers.AdaBelief
Optimisers.Lion
```
```