Flux chooses a third path, by encapsulating the parameters within the function.
The simplest way to do this is a closure, an anonymous function which Julia knows
to depend on some local variable θ3:

An equivalent, but tidier, way is to construct a struct in which to store the parameters.
Any struct can be made callable, allowing its instances to act just like function:

Internally, there is little difference between a closure and a struct.
They have the same fields, and equivalent methods:

The virtue of encapsulation is that it makes composition very easy.
We can make more complicated functions by combining simple ones,
and each will keep track of its own parameters.
Juia writes function composition as ∘, for instance (inv ∘ sin)(pi/6) ≈ 2,
and we can use exactly this for our parameterised polynomials:

Flux models are precisely made by such function composition.
In fact, poly3  and poly4 are already valid Flux models.

Structural Gradients Structural Gradients')" title="Copy chapter prompt for LLMs"> Copy Chapter
The derivative of a scalar function is its slope: how fast the output changes as the input is changed slightly.
This may be found approximately by evaluating at two nearby points, and exactly by taking the limit in
which the distance between them approaches zero:

Flux's gradient(f, x) works this out for f(x), and gives exactly ∂f/∂x = 2.0 here:

The reason gradient returns a tuple, not just the number 2.0, is to allow for
functions taking several arguments. (That's also why it's not called "derivative".)
For instance, this returns ∂f/∂x, ∂f/∂y, ∂f/∂z:

For our parameterised polynomial, we have ∂f/∂x but we are really more interested
in ∂f/∂θ, as this will tell us about how the parameters are affecting the answer.
It is not impossible to track gradients with respect to global θ, but much clearer to track explicit arguments.
Here's how this works for poly2 (which takes θ as a 2nd argument) and poly3 (which encapsulates θ):

The first entry is ∂f/∂x as before, but the second entry is more interesting.
For poly2, we get ∂f/∂θ as grad2[2] directly.
It is a vector, because θ is a vector, and has elements [∂f/∂θ[1], ∂f/∂θ[2], ∂f/∂θ[3]].
For poly3s, however, we get a NamedTuple whose fields correspond to those of the struct Poly3.
This is called a structural gradient. And the nice thing about them is that they work for
arbitrarily complicated structures, for instance:

Here grad4[2].inner.θ3 corresponds to poly4.inner.θ3.
These matching nested structures are at the core of how Flux works.
!!! note "Implicit gradients"
    Earlier versions of Flux used a different way to relate parameters and gradients,
    which looks like this:

Here Params is a set of references to global variables using objectid,
    and g1 isa Grads is a dictionary from these to their gradients.
    This method of gradient takes a zero-argument function, which only implicitly
    depends on θ.
Automatic Differentiation  Copy Chapter
Flux's gradient function by default calls a companion packages called Zygote.
Zygote performs source-to-source automatic differentiation, meaning that gradient(f, x)
hooks into Julia's compiler to find out what operations f contains, and transforms this
to produce code for computing ∂f/∂x.
Zygote can in principle differentiate almost any Julia code. However, it's not perfect,
and you may eventually want to read its page about limitations.
In particular, a major limitation is that mutating an array is not allowed.
Flux can also be used with other automatic differentiation (AD) packages.
It was originally written using Tracker, a more traditional operator-overloading approach.
The future might be Enzyme, and Flux now builds in an easy way to use this instead, turned on by wrapping the model in Duplicated. (For details, see the Enzyme page in the manual.)

Flux.gradient follows Zygote's convention that arguments with no derivative are marked nothing.
Here, this is because Const(5.0) is explicitly constant.
Below, we will see an example where nothing shows up because the model struct has fields containing things other than parameters, such as an activation function.
(It also adopts the convention that gradient(f, x, y) returns a tuple (∂f/∂x, ∂f/∂y), without a "∂f/∂f" term for the function. This is why we had to write gradient(|>, 5, poly4) above, not just gradient(poly4, 5).)
The function withgradient works the same way, but also returns the value of the function:

One can also directly specify which AD backend to use, by passing an adtype among the supported ones
(AutoMooncake, AutoEnzyme, AutoZygote, AutoFiniteDifferences) as the second argument.
The corresponding AD package has to be loaded first.
Here is an example using Mooncake:

and here is the same example using Enzyme:

Simple Neural Networks  Copy Chapter
The polynomial functions above send a number x to another a number y.
Neural networks typically take a vector of numbers, mix them all up, and return another vector.
Here's a very simple one, which will take a vector like x = [1.0, 2.0, 3.0]
and return another vector y = layer1(x) with length(y) == 2:

Here sigmoid is a nonlinear function, applied element-wise
because it is called with .(), called broadcasting.
Like poly1 above, this layer1 has as its parameters the global variables W, b.
We can similarly define a version which takes these as arguments (like poly2),
and a version which encapsulates them (like poly3 above):

This third way is precisely a Flux model. And we can again make a tidier version
using a struct to hold the parameters:

The one new thing here is a friendly constructor Layer(in, out, act).
This is because we anticipate composing several instances of this thing,
with independent parameter arrays, of different sizes and different
random initial parameters.
Let's try this out, and look at its gradient:

This ∂f/∂layer3s is a named tuple with the same fields as Layer.
Within it, the gradient with respect to W is a matrix of seemingly random numbers.
Notice that there is also an entry for act, which is nothing,
as this field of the struct is not a smoothly adjustable parameter.
We can compose these layers just as we did the polynomials above, in poly4.
Here's a composition of 3 functions, in which the last step is the function only
which takes a 1-element vector and gives us the number inside:

This gradient is starting to be a complicated nested structure.
But it works just like before: grad.outer.inner.W corresponds to model1.outer.inner.W.
We don't have to use ∘ (which makes a ComposedFunction struct) to combine layers.
Instead, we could define our own container struct, or use a closure.
This model2 will work the same way (although its fields have different names):

Rather than define everything from scratch every time, Flux provides a library of
commonly used layers. The same model could be defined:

How does this model3 differ from the model1 we had before?
* Flux's Chain works left-to-right, the reverse of Base's ∘.
  Its contents is stored in a tuple, thus model3.layers[1].weight is an array.
* Flux's layer Dense has only minor differences from our struct Layer:
  - Like struct Poly3{T} above, it has type parameters for its fields -- the compiler does not know exactly what type layer3s.W will be, which costs speed.
  - Its initialisation uses not randn (normal distribution) but glorot_uniform by default.
  - It reshapes some inputs (to allow several batch dimensions), and produces more friendly errors on wrong-size input.
  - And it has some performance tricks: making sure element types match, and re-using some memory.
* The function σ is calculated in a slightly better way,
  and has a rule telling Zygote how to differentiate it efficiently.
* Flux overloads Base.show so to give pretty printing at the REPL prompt.
  Calling Flux.@layer Layer will add this, and some other niceties.
All Flux layers accept a batch of samples: Instead of mapping one sample x::Vector to one output y::Vector, they map columns of a matrix xs::Matrix to columns of the output. This looks like f(xs) ≈ stack(f(x) for x in eachcol(xs)) but is done more efficiently.
If what you need isn't covered by Flux's built-in layers, it's easy to write your own.
There are more details later, but the steps are invariably those shown for struct Layer above:
1. Define a struct which will hold the parameters.
2. Make it callable, to define how it uses them to transform the input x
3. Define a constructor which initialises the parameters (if the default constructor doesn't do what you want).
4. Annotate with @layer to opt-in to pretty printing, and other enhancements.

To deal with such nested structures, Flux relies heavily on an associated package
called Functors. Its basic function is fmap,
which generalises map(f, x) to work on almost anything.
For example, this is how gpu moves all arrays within a model to the GPU,
reconstructing another only ∘ Layer(...) ∘ Layer(...) (or a Chain etc.) around the new CuArrays:

And this is a very simple gradient update of the parameters, walking over model and grad simultaneously:

!!! note
    Before Flux v0.15 (and Functors v0.5), this exploration of structs was opt-in.
    After defining struct Layer it was necessary to call @functor Layer (or @layer Layer) before Flux would look inside.
    This has now changed to be opt-out: Functors (and hence Flux) will explore arbitrary structs, unless told not to (using Functors.@leaf).
    This is why even "anonymous structs" created by closures, like poly3 and layer3 above, are now valid Flux models, although the use of named structs is still recommended practice.
Curve Fitting  Copy Chapter
Above we took gradients of the output, or sometimes to the first element
of the output -- it must be a number, not a vector. Adjusting the parameters
to make this smaller won't lead us anywhere interesting. Instead, we should minimise
some loss function which compares the actual output to our desired output.
Perhaps the simplest example is curve fitting. The previous page fitted
a linear model to data. With our two-layer model, we can fit a nonlinear function.
For example, let us use f(x) = 2x - x^3 evaluated at some points x in -2:0.1:2 as the data,
and adjust the parameters of model3 from above so that its output is similar.

The same code will also work with model1 or model2 instead.
Here's how to plot the desired and actual outputs:

More detail about what exactly the function train! is doing, and how to use rules other than simple Descent, is what the next page in this guide is about: training.
---
Src/Guide/Models/Overview  Copy Section
Flux Overview: Fitting a Straight Line
Flux is a pure Julia ML stack that allows you to build predictive models. Here are the steps for a typical Flux program:
1. Provide training and test data
2. Build a model with configurable parameters to make predictions
3. Iteratively train the model by tweaking the parameters to improve predictions
4. Verify your model
Under the hood, Flux uses a technique called automatic differentiation to take gradients that help improve predictions. Flux is also fully written in Julia so you can easily replace any layer of Flux with your own code to improve your understanding or satisfy special requirements.
Here's how you'd use Flux to build and train the most basic of models, step by step.
A Trivial Prediction  Copy Section
This example will predict the output of the function 4x + 2. Making such predictions is called "linear regression", and is really too simple to need a neural network. But it's a nice toy example.
First, import Flux and define the function we want to simulate:

This example will build a model to approximate the actual function.
1. Provide Training and Test Data  Copy Chapter
Use the actual function to build sets of data for training and verification:

Normally, your training and test data come from real world observations, but here we simulate them.
2. Build a Model to Make Predictions  Copy Chapter
Now, build a model to make predictions with 1 input and 1 output:

Under the hood, a dense layer is a struct with fields weight and bias. weight represents a weights' matrix and bias represents a bias vector. There's another way to think about a model. In Flux, models are conceptually predictive functions:

Dense(1 => 1) also implements the function σ(Wx+b) where W and b are the weights and biases. σ is an activation function (more on activations later). Our model has one weight and one bias, but typical models will have many more. Think of weights and biases as knobs and levers Flux can use to tune predictions. Activation functions are transformations that tailor models to your needs. 
This model will already make predictions, though not accurate ones yet:

In order to make better predictions, you'll need to provide a loss function to tell Flux how to objectively evaluate the quality of a prediction. Loss functions compute the cumulative distance between actual values and predictions.

More accurate predictions will yield a lower loss. You can write your own loss functions or rely on those already provided by Flux. This loss function is called mean squared error (and built-in as mse). Flux works by iteratively reducing the loss through training.
3. Improve the Prediction  Copy Chapter
Under the hood, the Flux Flux.train! function uses a loss function and training data to improve the parameters of your model based on a pluggable optimiser:

Now, we have the optimiser and data we'll pass to train!. All that remains are the parameters of the model. Remember, each model is a Julia struct with a function and configurable parameters. Remember, the dense layer has weights and biases that depend on the dimensions of the inputs and outputs:

The dimensions of these model parameters depend on the number of inputs and outputs.
Flux will adjust predictions by iteratively changing these parameters according to the optimiser.
This optimiser implements the classic gradient descent strategy. Now improve the parameters of the model with a call to Flux.train! like this:

And check the loss:

It went down. Why?

The parameters have changed. This single step is the essence of machine learning.
3+. Iteratively Train the Model  Copy Chapter
In the previous section, we made a single call to train! which iterates over the data we passed in just once. An epoch refers to one pass over the dataset. Typically, we will run the training for multiple epochs to drive the loss down even further. Let's run it a few more times:

After 200 training steps, the loss went down, and the parameters are getting close to those in the function the model is built to predict.
4. Verify the Results  Copy Chapter
Now, let's verify the predictions:

The predictions are good. Here's how we got there. 
First, we gathered real-world data into the variables x_train, y_train, x_test, and y_test. The x_ data defines inputs, and the y_ data defines outputs. The _train data is for training the model, and the _test data is for verifying the model. Our data was based on the function 4x + 2.
Then, we built a single input, single output predictive model, predict = Dense(1 => 1). The initial predictions weren't accurate, because we had not trained the model yet.
After building the model, we trained it with train!(loss, predict, data, opt). The loss function is first, followed by the model itself, the training data, and the Descent optimiser provided by Flux. We ran the training step once, and observed that the parameters changed and the loss went down. Then, we ran the train! many times to finish the training process.
After we trained the model, we verified it with the test data to verify the results. 
This overall flow represents how Flux works. Let's drill down a bit to understand what's going on inside the individual layers of Flux.
---
Src/Guide/Models/Quickstart  Copy Section
A Neural Network in One Minute
If you have used neural networks before, then this simple example might be helpful for seeing how the major parts of Flux work together. Try pasting the code into the REPL prompt.
If you haven't, then you might prefer the Fitting a Straight Line page.

Here's the loss during training:

This XOR ("exclusive or") problem is a variant of the famous one which drove Minsky and Papert to invent deep neural networks in 1969. For small values of "deep" -- this has one hidden layer, while earlier perceptrons had none. (What they call a hidden layer, Flux calls the output of the first layer, model1.)
Since then things have developed a little. 
Features to Note  Copy Chapter
Some things to notice in this example are:
* The batch dimension of data is always the last one. Thus a 2×1000 Matrix is a thousand observations, each a column of length 2. Flux defaults to Float32, but most of Julia to Float64.
* The model can be called like a function, y = model(x). Each layer like Dense is an ordinary struct, which encapsulates some arrays of parameters (and possibly other state, as for BatchNorm).
* But the model does not contain the loss function, nor the optimisation rule. The momenta needed by Adam are stored in the object returned by setup. And Flux.logitcrossentropy is an ordinary function that combines the softmax and crossentropy functions.
* The do block creates an anonymous function, as the first argument of gradient. Anything executed within this is differentiated.
Instead of calling gradient and update! separately, there is a convenience function train!. If we didn't want anything extra (like logging the loss), we could replace the training loop with the following:

If you want to keep control of the loop — for logging, early stopping, or a custom schedule — but would rather not spell out gradient and update! yourself, trainstep! is the per-step primitive that train! is built on. It evaluates the loss, computes the gradient, and updates model and opt_state in place, returning the loss. The batch is passed as a tuple and spliced into the loss after the model (so it is called as loss(m, x, y)):

This is exactly the original training loop above, with the withgradient/update! pair collapsed into a single call. Use trainstep_withgradient! if you also need the gradient. On a Reactant device, trainstep! (and train!) additionally compile and cache the whole step into one executable.
* Notice that the full dataset noisy lives on the CPU, and is moved to the GPU one batch at a time, by xy_cpu |> device. This is generally what you want for large datasets. Calling loader |> device similarly modifies the DataLoader to move one batch at a time.
* In our simple example, we conveniently created the model has a Chain of layers.
For more complex models, you can define a custom struct MyModel containing layers and arrays and implement the call operator (::MyModel)(x) = ... to define the forward pass. This is all that is needed for Flux to work. Marking the struct with Flux.@layer will add some more functionality, like pretty printing and the ability to mark some internal fields as trainable or not (also see trainable).
---
Src/Guide/Models/Recurrence  Copy Section
Recurrent Models
Recurrent cells  Copy Chapter
To introduce Flux's recurrence functionalities, we will consider the following vanilla recurrent neural network structure:
In the above, we have a sequence of length 3, where x1 to x3 represent the input at each step. It could be a timestamp or a word in a sentence encoded as vectors. y1 to y3 are their respective outputs.
An aspect to recognise is that in such a model, the recurrent cells A all refer to the same structure. What distinguishes it from a simple dense layer is that the cell A is fed, in addition to an input x, with information from the previous state of the model (hidden state denoted as h1 & h2 in the diagram).
In the most basic RNN case, cell A could be defined by the following:

Notice how the above is essentially a Dense layer that acts on two inputs, xt and ht.
The result of the forward pass at each time step, is a tuple containing the output yt and the updated state ht. The updated state is used as an input in next iteration. In the simple case of a vanilla RNN, the 
output and the state are the same. In more complex cells, such as LSTMCell, the state can contain multiple arrays.
There are various recurrent cells available in Flux, notably RNNCell, LSTMCell and GRUCell, which are documented in the layer reference. The hand-written example above can be replaced with:

The entire output y or just the last output y[end] can be used for further processing, such as classification or regression. 
Using a cell as part of a model  Copy Chapter
Let's consider a simple model that is trained to predict a scalar quantity for each time step in a sequence. The model will have a single RNN cell, followed by a dense layer to produce the output.
Since the RNNCell can deal with batches of data, we can define the model to accept an input where
at each time step, the input is a matrix of size (input_size, batch_size).

Notice that we stack the hidden states z to form a tensor of size (hidden_size, seq_len, batch_size). This can speed up the final classification, since we then process all the outputs at once with a single forward pass of the dense layer. 
Let's now define the training loop for this model:

Handling the whole sequence at once  Copy Chapter
In the above example, we processed the sequence one time step at a time using a recurrent cell. However, it is possible to process the entire sequence at once. This can be done by stacking the input data x to form a tensor of size (input_size, seq_len) or (input_size, seq_len, batch_size). 
One can then use the RNN, LSTM or GRU layers to process the entire input tensor. 
Let's consider the same example as above, but this time we use an RNN layer instead of an RNNCell:

Finally, the Recurrence layer can be used wrap any recurrent cell to process the entire sequence at once. For instance, a type behaving the same as the LSTM layer can be defined as follows:

Stacking recurrent layers  Copy Chapter
Recurrent layers can be stacked to form a deeper model by simply chaining them together using the Chain layer. The output of a layer is fed as input to the next layer in the chain.
For instance, a model with two LSTM layers can be defined as follows:

If more fine grained control is needed, for instance to have a trainable initial hidden state, one can define a custom model as follows:

---
Src/Guide/Training/Training  Copy Section
Training a Flux Model
Training refers to the process of slowly adjusting the parameters of a model to make it work better.
Besides the model itself, we will need three things:
 An objective function* that evaluates how well a model is doing on some input.
 An optimisation rule* which describes how the model's parameters should be adjusted.
 Some training data* to use as the input during this process.
Usually the training data is some collection of examples (or batches of examples) which
are handled one-by-one. One epoch of training means that each example is used once,
something like this:

This loop can also be written using the function train!,
but it's helpful to understand the pieces first:

Model Gradients  Copy Chapter
Fist recall from the section on taking gradients that
Flux.gradient(f, a, b) always calls f(a, b), and returns a tuple (∂f_∂a, ∂f_∂b).
In the code above, the function f passed to gradient is an anonymous function with
one argument, created by the do block, hence  grads is a tuple with one element.
Instead of a do block, we could have written:

Since the model is some nested set of layers, grads[1] is a similarly nested set of
NamedTuples, ultimately containing gradient components. If (for example) 
θ = model.layers[1].weight[2,3] is one scalar parameter, an entry in a matrix of weights,
then the derivative of the loss with respect to it is ∂f_∂θ = grads[1].layers[1].weight[2,3].
It is important that the execution of the model takes place inside the call to gradient,
in order for the influence of the model's parameters to be observed by Zygote.
It is also important that every update! step receives a newly computed gradient,
as it will change whenever the model's parameters are changed, and for each new data point.

Loss Functions  Copy Chapter
The objective function must return a number representing how far the model is from
the desired result. This is termed the loss of the model.
This number can be produced by any ordinary Julia code, but this must be executed
within the call to gradient. For instance, we could define a function

or write this directly inside the do block above. Many commonly used functions,
like mse for mean-squared error or crossentropy for cross-entropy loss,
are available from the Flux.Losses module.

Optimisation Rules  Copy Chapter
The simplest kind of optimisation using the gradient is termed gradient descent
(or sometimes stochastic gradient descent when, as here, it is not applied to the entire dataset at once).
Gradient descent needs a learning rate which is a small number describing how fast to walk downhill,
usually written as the Greek letter "eta", η. This is often described as a hyperparameter,
to distinguish it from the parameters which are being updated θ = θ - η * ∂loss_∂θ.
We want to update all the parameters in the model, like this:

A slightly more refined version of this loop to update all the parameters is wrapped up as a function update!(opt_state, model, grads[1]).
And the learning rate is the only thing stored in the Descent struct.
However, there are many other optimisation rules, which adjust the step size and
direction in various clever ways.
Most require some memory of the gradients from earlier steps, rather than always
walking straight downhill -- Momentum is the simplest.
The function setup creates the necessary storage for this, for a particular model.
It should be called once, before training, and returns a tree-like object which is the
first argument of update!. Like this:

Many commonly-used optimisation rules, such as Adam, are built-in.
These are listed on the optimisers page.
!!! compat "Implicit-style optimiser state"
    This setup makes another tree-like structure. Old versions of Flux did not do this,
    and instead stored a dictionary-like structure within the optimiser Adam(0.001).
    This was initialised on first use of the version of update! for "implicit" parameters.

Datasets & Batches  Copy Chapter
The loop above iterates through train_set, expecting at each step a tuple (input, label).
The very simplest such object is a vector of tuples, such as this:

or data = [(x, y), (x, y), (x, y)] for the same values three times.
Very often, the initial data is large arrays which you need to slice into examples.
To produce one iterator of pairs (x, y), you might want zip:

Here each iteration will use one matrix x (an image, perhaps) and one vector y.
It is very common to instead train on batches of such inputs (or mini-batches,
the two words mean the same thing) both for efficiency and for better results.
This can be easily done using the DataLoader:

Flux's layers are set up to accept such a batch of input data,
and the convolutional layers such as Conv require it.
The batch index is always the last dimension.
Training Loops  Copy Chapter
Simple training loops like the one above can be written compactly using
the train! function. Including setup, this reads:

Or explicitly writing the anonymous function which this do block creates,
train!((m,x,y) -> loss(m(x),y), model, train_set, opt_state) is exactly equivalent.
If you want to write the loop yourself but keep the convenience of a single gradient +
update call, trainstep! is the per-step primitive that
train! is built on. It differentiates the loss, updates model and opt_state in place,
and returns the loss:

Use trainstep_withgradient! if you also need the
gradient (it returns (loss, grad)). On a Reactant device both compile and cache the whole step,
just like train!.
Like withgradient, the loss passed to trainstep! may return auxiliary
data alongside the scalar loss — return a Tuple or NamedTuple whose first element is the loss,
and the gradient is taken of the loss alone while the whole value is returned. This is convenient
for logging a metric computed during the forward pass:

Real training loops often need more flexibility, and the best way to do this is just
to write the loop. This is ordinary Julia code, without any need to work through some
callback API. Here is an example, in which it may be helpful to note:
* The function withgradient is like gradient but also
  returns the value of the function, for logging or diagnostic use.
* Logging or printing is best done outside of the gradient call,
  as there is no need to differentiate these commands.
* To use result for logging purposes, you could change the do block to end with 
  return my_loss(result, label), result, i.e. make the function passed to withgradient
  return a tuple. The first element is always the loss.
* Julia's break and continue keywords let you exit from parts of the loop.

Regularisation  Copy Chapter
The term regularisation covers a wide variety of techniques aiming to improve the
result of training. This is often done to avoid overfitting.
Some of these can be implemented by simply modifying the loss function.
L₂ regularisation (sometimes called ridge regression) adds to the loss a penalty
proportional to θ^2 for every scalar parameter.
A very simple model could be implemented as follows:

Accessing each individual parameter array by hand won't work well for large models.
Instead, we can use Flux.trainables to collect all of them,
and then apply a function to each one, and sum the result:

However, the gradient of this penalty term is very simple: It is proportional to the original weights.
So there is a simpler way to implement exactly the same thing, by modifying the optimiser
instead of the loss function. This is done by replacing this:

with this:

Flux's optimisers are really modifications applied to the gradient before using it to update
the parameters, and OptimiserChain applies two such modifications.
The first, WeightDecay adds 0.42 times the original parameter to the gradient,
matching the gradient of the penalty above (with the same, unrealistically large, constant).
After that, in either case, Adam computes the final update.
The same trick works for L₁ regularisation (also called Lasso), where the penalty is 
pen_l1(x::AbstractArray) = sum(abs, x) instead. This is implemented by SignDecay(0.42).
The same OptimiserChain mechanism can be used for other purposes, such as gradient clipping with ClipGrad or ClipNorm.
Besides L1 / L2 / weight decay, another common and quite different kind of regularisation is
provided by the Dropout layer. This turns off some outputs of the
previous layer during training.
It should switch automatically, but see trainmode! / testmode! to manually enable or disable this layer.
Learning Rate Schedules  Copy Chapter
Finer control of training, you may wish to alter the learning rate mid-way through training.
This can be done with adjust!, like this:

Other hyper-parameters can also be adjusted, such as Flux.adjust!(opt_state, beta = (0.8, 0.99)).
And such modifications can be applied to just one part of the model.
For instance, this sets a different learning rate for the encoder and the decoder:

Scheduling Optimisers  Copy Chapter
In practice, it is fairly common to schedule the learning rate of an optimiser to obtain faster convergence. There are a variety of popular scheduling policies, and you can find implementations of them in ParameterSchedulers.jl. The documentation for ParameterSchedulers.jl provides a more detailed overview of the different scheduling policies, and how to use them with Flux optimisers. Below, we provide a brief snippet illustrating a cosine annealing schedule with a momentum optimiser.
First, we import ParameterSchedulers.jl and initialize a cosine annealing schedule to vary the learning rate between 1e-4 and 1e-2 every 10 epochs. We also create a new Momentum optimiser.

schedule can also be indexed (e.g. schedule(100)) or iterated like any iterator in Julia.
ParameterSchedulers.jl schedules are stateless (they don't store their iteration state). If you want a _stateful_ schedule, you can use ParameterSchedulers.Stateful:

Finally, a scheduling function can be incorporated into the optimser's state, advanced at each gradient update step, and possibly passed to the train! function. See this section of ParameterSchedulers.jl documentation for more details.
ParameterSchedulers.jl allows for many more scheduling policies including arbitrary functions, looping any function with a given period, or sequences of many schedules. See the ParameterSchedulers.jl documentation for more info.
Freezing layer parameters  Copy Chapter
To completely disable training of some part of the model, use freeze!.
This is a temporary modification, reversed by thaw!:

While adjust! and freeze!/thaw! make temporary modifications to the optimiser state,
permanently removing some fields of a new layer type from training is usually done
when defining the layer, by calling for example @layer NewLayer trainable=(weight,).
---
Src/Guide/Gpu  Copy Section
GPU Support
Most work on neural networks involves the use of GPUs, as they can typically perform the required computation much faster.
This page describes how Flux co-operates with various other packages, which talk to GPU hardware.
For those in a hurry, see the quickstart page. Or do using CUDA and then call gpu on both the model and the data. 
Basic GPU use: from Array to CuArray  to CuArray')" title="Copy chapter prompt for LLMs"> Copy Chapter
Julia's GPU packages work with special array types, in place of the built-in Array.
The most used is CuArray provided by CUDA.jl, for GPUs made by NVIDIA.
That package provides a function cu which converts an ordinary Array (living in CPu memory) to a CuArray (living in GPU memory).
Functions like * and broadcasting specialise so that, when given CuArrays, all the computation happens on the GPU:

Notice that cu doesn't only move arrays, it also recurses into many structures, such as the tuple (W, x) above.
(Notice also that it converts Julia's default Float64 numbers to Float32, as this is what most GPUs support efficiently -- it calls itself "opinionated". Flux defaults to Float32 in all cases.)
To use CUDA with Flux, you can simply use cu to move both the model, and the data.
It will create a copy of the Flux model, with all of its parameter arrays moved to the GPU:

Notice that you need using CUDA (every time) but also ] add cuDNN (once, when installing packages).
This is a quirk of how these packages are set up.
(The cuDNN.jl sub-package handles operations such as convolutions, called by Flux via NNlib.jl.)
Flux's gradient, and training functions like setup, update!, and train!, are all equally happy to accept GPU arrays and GPU models, and then perform all computations on the GPU.
It is recommended that you move the model to the GPU before calling setup.

To move arrays and other objects back to the CPU, Flux provides a function cpu.
This is recommended when saving models, Flux.state(c_model |> cpu), see below.

!!! compat "Flux ≤ 0.13"
    Old versions of Flux automatically loaded CUDA.jl to provide GPU support. Starting from Flux v0.14, it has to be loaded separately. Julia's package extensions) allow Flux to automatically load some GPU-specific code when needed.
Other GPU packages for AMD & Apple  Copy Chapter
Non-NVIDIA graphics cards are supported by other packages. Each provides its own function which behaves like cu.
AMD GPU support provided by AMDGPU.jl, on systems with ROCm and MIOpen installed.
This package has a function roc which converts Array to ROCArray:

Experimental support for Apple devices with M-series chips is provided by  Metal.jl. This has a function mtl which works like cu, converting Array to MtlArray:

!!! danger "Experimental"
    Metal support in Flux is experimental and many features are not yet available.
    AMD support is improving, but likely to have more rough edges than CUDA.
If you want your model to work with any brand of GPU, or none, then you may not wish to write cu everywhere.
One simple way to be generic is, at the top of the file, to un-comment one of several lines which import a package and assign its "adaptor" to the same name:

!!! note "Adapt.jl"
    The functions cu, mtl, roc all use Adapt.jl, to work within various wrappers.
    The reason they work on Flux models is that Flux.@layer Layer defines methods of Adapt.adapt_structure(to, lay::Layer).

Automatic GPU choice with gpu and gpu_device  and gpu_device')" title="Copy chapter prompt for LLMs"> Copy Chapter
Flux also provides a more automatic way of choosing which GPU (or none) to use. This is the function gpu:
* By default it does nothing.
* If the package CUDA is loaded, and CUDA.functional() === true, then it behaves like cu.
* If the package AMDGPU is loaded,  and AMDGPU.functional() === true, then it behaves like roc.
* If the package Metal is loaded, and Metal.functional() === true, then it behaves like mtl.
* If two different GPU packages are loaded, the first one takes priority.
For the most part, this means that a script which says model |> gpu and data |> gpu will just work.
It should always run, and if a GPU package is loaded (and finds the correct hardware) then that will be used.
The function gpu uses a lower-level function called gpu_device from MLDataDevices.jl,
which checks what to do and then returns some device object. In fact, the entire implementation is just this:

Automatic backend selection through gpu is not type-stable. That doesn't matter if you do it once, or once per large batch -- it costs a few microseconds. But it might matter if you do it within some loop.
To avoid this, you can first obtain a "device object" with device = gpu_device(), once, and then use this as the function to transfer data. Something like this:

Finally, setting a backend prefence with gpu_backend! gives type stability to the whole pipeline.
Transferring Training Data  Copy Chapter
In order to train the model using the GPU both model and the training data have to be transferred to GPU memory. Moving the data can be done in two different ways:
1. Iterating over the batches in a DataLoader object transferring each one of the training batches at a time to the GPU. This is recommended for large datasets. Done by hand, it might look like this:

Rather than write this out every time, you can just call gpu(::DataLoader):

This is equivalent to DataLoader(MLUtils.mapobs(gpu, (X, Y)); keywords...).
   Something similar can also be done with CUDA.CuIterator, gpu_train_loader = CUDA.CuIterator(train_loader). However, this only works with a limited number of data types: first(train_loader) should be a tuple (or NamedTuple) of arrays.
2. Transferring all training data to the GPU at once before creating the DataLoader. This is usually performed for smaller datasets which are sure to fit in the available GPU memory.

Here (X, Y) |> gpu applies gpu to both arrays, as it recurses into structures.
Saving GPU-Trained Models  Copy Chapter
After the training process is done, we must always transfer the trained model back to the CPU memory before serializing or saving to disk. This can be done with cpu:

and then

The reason behind this is that models trained in the GPU but not transferred to the CPU memory scope will expect CuArrays as input. In other words, Flux models expect input data coming from the same kind device in which they were trained on.
In controlled scenarios in which the data fed to the loaded models is guaranteed to be in the GPU there's no need to transfer them back to CPU memory scope, however in production environments, where artifacts are shared among different processes, equipments or configurations, there is no guarantee that the CUDA.jl package will be available for the process performing inference on the model loaded from the disk.

Disabling CUDA or choosing which GPUs are visible to Flux  Copy Chapter
Sometimes it is required to control which GPUs are visible to julia on a system with multiple GPUs or disable GPUs entirely. This can be achieved with an environment variable CUDA_VISIBLE_DEVICES.
To disable all devices:

To select specific devices by device id:

More information for conditional use of GPUs in CUDA.jl can be found in its documentation, and information about the specific use of the variable is described in the Nvidia CUDA blog post.

Data movement across GPU devices  Copy Chapter
Flux also supports getting handles to specific GPU devices, and transferring models from one GPU device to another GPU device from the same backend. Let's try it out for NVIDIA GPUs. First, we list all the available devices:

Then, let's select the device with id 0:

Notice that indexing starts from 0 in the CUDA.devices() output, but gpu_device! expects the device id starting from 1.
Then, let's move a simple dense layer to the GPU represented by device0:

Next, we'll get a handle to the device with id 1, and move dense_model to that device:

Due to a limitation in Metal.jl, currently this kind of data movement across devices is only supported for CUDA and AMDGPU backends.

Distributed data parallel training  Copy Chapter
!!! danger "Experimental"
    Distributed support is experimental and could change in the future.

Flux supports now distributed data parallel training with DistributedUtils module.
If you want to run your code on multiple GPUs, you have to install MPI.jl (see docs for more info).

Now you can run your code with mpiexecjl --project=. -n <np> julia <filename>.jl from CLI.
You can use either the MPIBackend or NCCLBackend, the latter only if also NCCL.jl is loaded. First, initialize a backend with DistributedUtils.initialize, e.g.

Pass your model, as well as any data to GPU device.

In this case, we are training on a total of 16 * number of processes samples. You can also use DistributedUtils.DistributedDataContainer to split the data uniformly across processes (or do it manually).

You have to wrap your model in DistributedUtils.FluxDistributedModel and synchronize it (broadcast across all processes):

Time to set up an optimizer by using DistributedUtils.DistributedOptimizer and synchronize it as well.

Now you can define loss and train the model.

Remember that in order to run it on multiple GPUs you have to run from CLI mpiexecjl --project=. -n <np> julia <filename>.jl,
where  <np> is the number of processes that you want to use. The number of processes usually corresponds to the number of gpus.
By default MPI.jl MPI installation is CUDA-unaware so if you want to run it in CUDA-aware mode, read more here on custom installation and rebuilding MPI.jl.
Then test if your MPI is CUDA-aware by

If it is, set your local preference as below

!!! warning "Known shortcomings"
    We don't run CUDA-aware tests so you're running it at own risk.

Checking GPU Availability  Copy Chapter
By default, Flux will run the checks on your system to see if it can support GPU functionality. You can check if Flux identified a valid GPU setup by typing the following:

For AMD GPU:

For Metal GPU:

---
Src/Guide/Performance  Copy Section
Performance Tips
All the usual Julia performance tips apply.
As always profiling your code is generally a useful way of finding bottlenecks.
Below follow some Flux specific tips/reminders.
Don't use more precision than you need  Copy Chapter
Flux works great with all kinds of number types.
But often you do not need to be working with say Float64 (let alone BigFloat).
Switching to Float32 can give you a significant speed up,
not because the operations are faster, but because the memory usage is halved.
Which means allocations occur much faster.
And you use less memory.
Preserve inputs' types  Copy Chapter
Not only should your activation and loss functions be type-stable,
they should also preserve the type of their inputs.
A very artificial example using an activation function like

will result in performance on Float32 input orders of magnitude slower than the normal tanh would,
because it results in having to use slow mixed type multiplication in the dense layers.
Similar situations can occur in the loss function during backpropagation.
Which means if you change your data say from Float64 to Float32 (which should give a speedup: see above),
you will see a large slow-down.
This can occur sneakily, because you can cause type-promotion by interacting with a numeric literals.
E.g. the following will have run into the same problem as above:

While one could change the activation function (e.g. to use 0.01f0*x), the idiomatic (and safe way)  to avoid type casts whenever inputs changes is to use oftype:

Evaluate batches as matrices of features  Copy Chapter
While it can sometimes be tempting to process your observations (feature vectors) one at a time
e.g.

It is much faster to concatenate them into a matrix,
as this will hit BLAS matrix-matrix multiplication, which is much faster than the equivalent sequence of matrix-vector multiplications.
The improvement is enough that it is worthwhile allocating new memory to store them contiguously.

When doing this kind of concatenation use reduce(hcat, xs) rather than hcat(xs...).
This will avoid the splatting penalty, and will hit the optimised reduce method.
Be aware of GPU memory inefficiencies  Copy Chapter
Currently, GPU memory is not handled as well as system memory. If your training loop is allocating significantly on the GPU, you can quickly fill your GPU memory and the piecemeal reclamation and shuffling of data between GPU and system memory can become extremely slow. If profiling shows that a significant portion of time is spent in the gpu function and your data sizes are not large, this may be the cause. Running an incremental garbage collection manually (GC.gc(false)) at regular intervals can keep your GPU memory free and responsive. See other tips for CUDA memory management here.
Compile your training with Reactant  Copy Chapter
The tips above tune eager execution. For the largest speedups, compile your model with Reactant.jl: it traces your code and lowers it — through MLIR and XLA — into a single fused, kernel-optimised executable that is cached and reused across calls. For training, the forward pass, the Enzyme reverse pass and the optimiser update are compiled together, which removes per-operation dispatch and kernel-launch overhead and lets the compiler fuse operations and plan buffers across the whole step.
You do not have to invoke the compiler yourself: when the model lives on a Reactant device, trainstep! — and train!, which is built on it — compile and cache the fused step automatically.

The cached executable is keyed on the batch shape, so keeping batch sizes fixed avoids recompilation (a smaller final batch simply compiles one extra executable). The first call pays a one-time compilation cost that can be large, so this pays off over many steps rather than a handful. See Compiling Flux with Reactant for the full guide, including inference and manual @compile.
---
Src/Guide/Reactant  Copy Section
Compiling Flux with Reactant
Reactant.jl traces your Julia code and compiles it — through MLIR and XLA — into a single optimised executable that runs on CPU, NVIDIA/AMD GPUs, or TPUs. For Flux this means the whole forward pass (and, for training, the Enzyme reverse pass and the optimiser update) is fused, kernel-optimised, and reused across calls, which is often substantially faster and more memory-efficient than eager execution.
This guide builds up in three steps:
1. A manually compiled example, so you can see exactly what Reactant does.
2. The trainstep! API, which compiles and caches a single training step for you.
3. train!, the full training loop, which is built on top of trainstep!.
Installation  Copy Chapter
Reactant is a normal dependency — add it and load it alongside Flux:

Reactant automatically selects a GPU backend if one is available, falling back to the CPU otherwise. See the Reactant GPU configuration docs for how to control this.
Moving data and models to the device  Copy Chapter
Just like gpu/cpu, Flux provides a device object that moves a model (or any nested structure of arrays) onto the Reactant device, converting its arrays into Reactant's ConcreteRArrays:

As with the GPU adaptors, dev recurses into structures, so (x, y) |> dev moves both, and it uses Float32 by default. Everything you feed to a compiled function must already live on the device.
1. A manually compiled example  Copy Chapter
Reactant does not run your code eagerly. Instead you compile a function once for a given set of input shapes and types, and then call the returned executable. The two entry points are:
- Reactant.@compile f(args...) — trace and compile f, returning a callable executable. Call it later with device-resident arguments of the same shapes.
- Reactant.@jit f(args...) — compile and immediately run, a convenient shortcut for one-off calls.
Compiling the forward pass (inference)  Copy Section

Note the call is forward(x), not forward(model, x): in @compile model(x) the model is the callee, so Reactant captures it (and its parameter arrays) inside the compiled executable, and you pass only the remaining arguments. Because the model is captured by reference, mutating its parameters in place — as training does — is reflected the next time you call forward. (Contrast this with the training-step examples below, where the model sits in an argument position, e.g. @compile my_step!(loss, model, x, y, opt_state), and so must be passed at call time.)
A compiled executable is specialised to the shapes of its inputs. Calling forward with a differently-shaped x (e.g. a smaller final batch) requires compiling a separate executable for that shape.
!!! note "trainmode / testmode"
    Layers such as Dropout and BatchNorm behave differently during training and inference. Call testmode!(model) before compiling an inference function, and trainmode!(model) before a training step, exactly as you would without Reactant.
Compiling a training step manually  Copy Section
To differentiate under Reactant, use Flux's withgradient with the AutoEnzyme backend — Reactant relies on Enzyme for AD. You can compile the value-and-gradient computation, or a whole step that also updates the model in place:

Reactant traces the mutation of model and opt_state and fuses the whole step. Writing this by hand gives you full control, but it is boilerplate that Flux can handle for you — which is what the next section is about.
2. The trainstep! API  API')" title="Copy chapter prompt for LLMs"> Copy Chapter
trainstep! performs exactly the step above — differentiate the loss, update model and opt_state in place, return the loss — and when the model lives on a Reactant device it compiles and caches the fused step automatically. You do not write @compile yourself, and repeated calls with the same model, optimiser, loss and batch shape reuse the cached executable.

The batch is passed as a tuple (x, y) and spliced into the loss as loss(model, x, y). The returned loss is read back to the host as an ordinary number.
Use trainstep_withgradient! if you also need the gradient — it returns (loss, grad), with the gradient left on the device. (Returning the gradient makes it an output of the compiled step and raises peak memory, so prefer trainstep! when you don't need it.)
Auxiliary loss outputs  Copy Section
Like withgradient, the loss may return auxiliary data alongside the scalar loss: return a Tuple or NamedTuple whose first element is the loss, and the gradient is taken of the loss alone while the whole value is returned. This is convenient for logging a metric computed in the forward pass — and, on Reactant, the metric is computed on-device as part of the same compiled forward:

3. The train! API  API')" title="Copy chapter prompt for LLMs"> Copy Chapter
train! is a loop over the data built on top of trainstep!, so on a Reactant device it inherits the same automatic compile-and-cache. Move the model to the device and call setup first, and provide device-resident data — train! rejects host arrays on the Reactant path.

train! runs a single step per batch, shows a progress bar, and stops with a DomainError if the loss becomes non-finite. Because the compiled step is cached and keyed on the batch shape, multi-epoch training does not recompile, and a smaller final batch simply compiles one additional executable that is then reused.
Tips and gotchas  Copy Chapter
- setup after moving to the device. Always model |> reactant_device() before Flux.setup, so the optimiser state is created from Reactant arrays (this keeps stateful rules like Adam's bias-correction term on the device).
- Device-resident data. Every input to a compiled step must already live on the device. Move batches with |> reactant_device(); a common pattern is to wrap a DataLoader with the device so batches are moved lazily.
- Compilation is keyed on shape. The first call for each distinct batch shape compiles; subsequent calls reuse the executable. Keeping batch sizes fixed (or accepting one extra compile for a smaller final batch) avoids repeated compilation.
- Watch the compile cache. Flux caches one executable per distinct (model, optimiser, loss, batch-shape). If you rebuild the model or optimiser every iteration, each iteration compiles a fresh step — Flux warns once the cache grows past a handful of entries. Entries are freed automatically when their model is garbage-collected.
- Reading results back. trainstep!/train! return host-side numbers already. For a manually compiled function, move device arrays back with cpu (or Reactant.to_number for a scalar).
- AD backend. On a Reactant device training differentiates with Enzyme by default. Passing an explicit adtype such as AutoZygote() or AutoMooncake() compilation is likely to fail. Duplicated models are not supported here — pass the plain model that already lives on the device.
See also the resnet_tinyimagenet example for a complete Reactant training script, including compiling a separate evaluation executable per batch shape.
---
Src/Guide/Saving  Copy Section
Saving and Loading Models
You may wish to save models so that they can be loaded and run in a later
session. Flux provides a number of ways to do this. 
The recommended way, which is the most robust one for long term storage, 
is to use Flux.state in combination with a serialization format like
JLD2.jl or
BSON.jl.
Save a model:

Load it again in a new session using Flux.loadmodel!:

!!! note
    If a saved model's parameters are stored on the GPU, the model will not load
    later on if there is no GPU support available. It's best to move your model
    to the CPU with cpu(model) before saving it.

Checkpointing  Copy Chapter
In longer training runs it's a good idea to periodically save your model, so that you can resume if training is interrupted (for example, if there's a power cut).

This will update the "model-checkpoint.jld2" every epoch.
You can get more advanced by saving a series of models throughout training, for example

will produce a series of models like "model-2018-03-06T02:57:10.41.jld2". You
could also store the current test set loss, so that it's easy to (for example)
revert to an older copy of the model if it starts to overfit.

Note that to resume a model's training, you might need to restore other stateful parts of your training loop. Possible examples are the optimiser state and the randomness used to partition the original data into the training and validation sets.
You can store the optimiser state alongside the model, to resume training
exactly where you left off:

Saving Models as Julia Structs
Models are just normal Julia structs, so it's fine to use any Julia storage
format to save the struct as it is instead of saving the state returned by Flux.state. 
BSON.jl is particularly convenient for this,
since it can also save anonymous functions, which are sometimes part of a model definition.
Save a model:

Load it again in a new session:

!!! warning
    Saving models this way could lead to compatibility issues across julia versions
    and across Flux versions if some of the Flux layers' internals are changed.
    It is therefore not recommended for long term storage, use Flux.state instead.
---
Src/Reference/Data/Mldatadevices  Copy Section

Transferring data across devices
Flux relies on the MLDataDevices.jl package to manage devices and transfer data across them. You don't have to explicitly use the package, as Flux re-exports the necessary functions and types.

---
Src/Reference/Data/Mlutils  Copy Section

Working with Data, using MLUtils.jl
Flux re-exports the DataLoader type and utility functions for working with
data from MLUtils.
DataLoader ')" title="Copy chapter prompt for LLMs"> Copy Chapter
The DataLoader can be used to create mini-batches of data, in the format train! expects.

Utility Functions  Copy Chapter
The utility functions are meant to be used while working with data;
these functions help create inputs for your models or batch your dataset.

---
Src/Reference/Data/Onehot  Copy Section

One-Hot Encoding with OneHotArrays.jl
It's common to encode categorical variables (like true, false or cat, dog) in "one-of-k" or "one-hot" form. OneHotArrays.jl provides the onehot function to make this easy.

There is also a onecold function, which is an inverse of onehot. It can also be given an array of numbers instead of booleans, in which case it performs an argmax-like operation, returning the label with the highest corresponding weight.

For multiple samples at once, onehotbatch creates a batch (matrix) of one-hot vectors, and onecold treats matrices as batches.

Note that these operations returned OneHotVector and OneHotMatrix rather than Arrays. OneHotVectors behave like normal vectors but avoid any unnecessary cost compared to using an integer index directly. For example, multiplying a matrix with a one-hot vector simply slices out the relevant row of the matrix under the hood.
Function listing  Copy Chapter

---
Src/Reference/Models/Activation  Copy Section

Activation Functions from NNlib.jl
These non-linearities used between layers of your model are exported by the NNlib package.
Note that, unless otherwise stated, activation functions operate on scalars. To apply them to an array you can call σ.(xs), relu.(xs) and so on. Alternatively, they can be passed to a layer like Dense(784 => 1024, relu) which will handle this broadcasting.
Functions like softmax are sometimes described as activation functions, but not by Flux. They must see all the outputs, and hence cannot be broadcasted. See the next page for details.
Alphabetical Listing  Copy Chapter

One More  Copy Chapter
Julia's Base.Math also provides tanh, which can be used as an activation function.
Note that many Flux layers will automatically replace this with NNlib.tanh_fast when called, as Base's tanh is slow enough to sometimes be a bottleneck.

---
Src/Reference/Models/Functors  Copy Section

Recursive transformations from Functors.jl
Flux models are deeply nested structures, and Functors.jl provides tools needed to explore such objects, apply functions to the parameters they contain (e.g. for moving them to gpu), and re-build them.
!!! compat "Flux ≤ v0.14"
    All layers were previously defined with the Functors.@functor macro.
    This still works, but it is recommended that you use the new Flux.@layer macro instead.
    Both allow Flux.setup to see the parameters inside, and gpu to move them to the GPU, but Flux.@layer also overloads printing,
    and offers a way to define trainable at the same time.
!!! compat "Functors v0.5"
    With Functors.jl v0.5, which is required by Flux v0.15 and later, every custom type is a functor by default. This means that applying Flux.@layer to a type is no longer strictly necessary, but it is still recommended for addictional features like pretty-printing.
Functors.jl has its own notes on basic usage for more details. Additionally, the Advanced Model Building and Customisation page covers the use cases of Functors in greater details.

Moving models, or data, to the GPU  Copy Chapter
Flux provides some convenience functions based on fmap. Some (f16, bf16, f32, f64) change the precision of all arrays in a model. Others are used for moving a model to of from GPU memory:

---
Src/Reference/Models/Layers  Copy Section

Built-in Layer Types
If you started at the beginning of the guide, then you have already met the
basic Dense layer, and seen Chain for combining layers.
These core layers form the foundation of almost all neural networks.
The Dense exemplifies several features:
* It contains an an activation function, which is broadcasted over the output. Because this broadcast can be fused with other operations, doing so is more efficient than applying the activation function separately.
* It take an init keyword, which accepts a function acting like rand. That is, init(2,3,4) should create an array of this size. Flux has many such functions built-in. All make a CPU array, moved later with gpu if desired.
* The bias vector is always initialised Flux.zeros32. The keyword bias=false will turn this off, i.e. keeping the bias permanently zero.
* It is annotated with @layer, which means that Flux.setup will see the contents, and gpu will move their arrays to the GPU.
By contrast, Chain itself contains no parameters, but connects other layers together.
The section on dataflow layers introduces others like this.
Fully Connected  Copy Chapter

Perhaps Scale isn't quite fully connected, but it may be thought of as Dense(Diagonal(s.weights), s.bias), and LinearAlgebra's Diagonal is a matrix which just happens to contain many zeros.

Convolution Models  Copy Chapter
These layers are used to build convolutional neural networks (CNNs).
They all expect images in what is called WHCN order: a batch of 32 colour images, each 50 x 50 pixels, will have size(x) == (50, 50, 3, 32). A single grayscale image might instead have size(x) == (28, 28, 1, 1).
Besides images, 2D data, they also work with 1D data, where for instance stereo sound recording with 1000 samples might have size(x) == (1000, 2, 1). They will also work with 3D data, ndims(x) == 5, where again the last two dimensions are channel and batch.
To understand how strides and padding work, the article by Dumoulin & Visin has great illustrations.

MultiHeadAttention  Copy Chapter
The basic blocks needed to implement Transformer architectures. See also the functional counterparts
documented in NNlib's Attention section.

Pooling  Copy Section
These layers are commonly used after a convolution layer, and reduce the size of its output. They have no trainable parameters.

Upsampling  Copy Chapter
The opposite of pooling, these layers increase the size of an array. They have no trainable parameters.

Embedding Vectors  Copy Chapter
These layers accept an index, and return a vector (or several indices, and several vectors). The possible embedding vectors are learned parameters.

Dataflow Layers, or Containers Dataflow Layers, or Containers')" title="Copy chapter prompt for LLMs"> Copy Chapter
The basic Chain(F, G, H) applies the layers it contains in sequence, equivalent to H ∘ G ∘ F. Flux has some other layers which contain layers, but connect them up in a more complicated way: SkipConnection allows ResNet's residual connection.

Recurrent Models  Copy Chapter
Much like the core layers above, but can be used to process sequence data (as well as other kinds of structured data).

Normalisation & Regularisation  Copy Chapter
These layers don't affect the structure of the network but may improve training times or reduce overfitting. Some of them contain trainable parameters, while others do not.

Test vs. Train  Copy Section
Several normalisation layers behave differently under training and inference (testing). By default, Flux will automatically determine when a layer evaluation is part of training or inference.
!!! warning
    This automatic train/test detection works best with Zygote, the default
    automatic differentiation package. It may not work with other packages
    such as Tracker, Yota, or ForwardDiff.
The functions Flux.trainmode! and Flux.testmode! let you manually specify which behaviour you want. When called on a model, they will place all layers within the model into the specified mode.

---
Src/Reference/Models/Losses  Copy Section

Loss Functions
Flux provides a large number of common loss functions used for training machine learning models.
They are grouped together in the Flux.Losses module.
Loss functions for supervised learning typically expect as inputs a target y, and a prediction ŷ from your model.
In Flux's convention, the order of the arguments is the following

Most loss functions in Flux have an optional argument agg, denoting the type of aggregation performed over the
batch:

Function listing  Copy Chapter

---
Src/Reference/Models/Nnlib  Copy Section

Neural Network primitives from NNlib.jl
Flux re-exports all of the functions exported by the NNlib package. This includes activation functions, described on their own page. Many of the functions on this page exist primarily as the internal implementation of Flux layer, but can also be used independently.

Attention  Copy Chapter
Primitives for the MultiHeadAttention layer.

Softmax  Copy Chapter
Flux's Flux.logitcrossentropy uses NNlib.logsoftmax internally.

Pooling  Copy Chapter
Flux's AdaptiveMaxPool, AdaptiveMeanPool, GlobalMaxPool, GlobalMeanPool, 
MaxPool, and MeanPool use NNlib.PoolDims, NNlib.maxpool, and NNlib.meanpool as their backend.

Padding  Copy Chapter

Convolution  Copy Chapter
Flux's Conv and CrossCor layers use NNlib.DenseConvDims and NNlib.conv internally.

Dropout  Copy Chapter

Normalization  Copy Chapter
Flux's BatchNorm, InstanceNorm, GroupNorm, and LayerNorm layers wrap the functional normalization operators below as their backend.

Upsampling  Copy Chapter
Flux's Upsample layer uses NNlib.upsample_nearest, NNlib.upsample_bilinear, and NNlib.upsample_trilinear as its backend. Additionally, Flux's PixelShuffle layer uses NNlib.pixel_shuffle as its backend.

Batched Operations  Copy Chapter
Flux's Flux.Bilinear layer uses NNlib.batched_mul internally.

Gather and Scatter  Copy Chapter
Flux's Embedding layer uses NNlib.gather as its backend.

Sampling  Copy Chapter

Losses  Copy Chapter

Miscellaneous  Copy Chapter

---
Src/Reference/Training/Callbacks  Copy Section

Callback Helpers

Patience Helpers  Copy Chapter
Flux provides utilities for controlling your training procedure according to some monitored condition and a maximum patience. For example, you can use early_stopping to stop training when the model is converging or deteriorating, or you can use plateau to check if the model is stagnating.
For example, below we create a pseudo-loss function that decreases, bottoms out, and then increases. The early stopping trigger will break the loop before the loss increases too much.

The keyword argument distance of early_stopping is a function of the form distance(best_score, score). By default distance is - and init_score is Inf, meaning that the monitored metric f is expected to be decreasing and minimized. If you use a metric such that improvement is shown by increasing values (e.g. accuracy), you can customize the distance function and the init_score value to, for example, (best_score, score) -> score - best_score and -Inf, respectively.

early_stopping and plateau are both built on top of patience. You can use patience to build your own triggers that use a patient counter. For example, if you want to trigger when the loss is below a threshold for several consecutive iterations:

Both predicate in patience and f in early_stopping / plateau can accept extra arguments. You can pass such extra arguments to predicate or f through the returned function:

---
Src/Reference/Training/Gradients  Copy Section

Automatic Differentiation in Flux
Flux's gradient function uses Zygote by default, and also uses this function within train! to differentiate the model.
Zygote has its own documentation, in particular listing some important limitations.
Flux also has support for Enzyme.jl, documented below and for Mooncake.jl.

Generic Gradient Interface  Copy Chapter

Automatic Differentiation using Zygote.jl Automatic Differentiation using Zygote.jl')" title="Copy chapter prompt for LLMs"> Copy Chapter
The default AD backend in Flux is Zygote. Besides  gradient calculation, Zygote also supports
higher-order derivatives, Jacobians, Hessians, and pullbacks.

ChainRules for Zygote  Copy Chapter
Zygote uses ChainRules.jl to define how to differentiate functions.
Sometimes it is necessary to exclude some code, or a whole function, from automatic differentiation.
This can be done using the following methods:

To manually supply the gradient for one function, you should define a method of rrule. ChainRules has detailed documentation on how this works.

Gradient customization for other AD packages such as Enzyme and Mooncake has to be done according to their own documentation.
Automatic Differentiation using Enzyme.jl Automatic Differentiation using Enzyme.jl')" title="Copy chapter prompt for LLMs"> Copy Chapter
Enzyme.jl is a new package for automatic differentiation.
Like Zygote.jl, calling gradient(f, x) causes it to hooks into the compiler and transform code that is executed while calculating f(x), in order to produce code for ∂f/∂x.
But it does so much later in the optimisation process (on LLVM instead of Julia's untyped IR) which you can read about here].
It needs far fewer custom rules than Zygote/ChainRules, and in particular is able to support mutation of arrays.
Flux now builds in support for this, using Enzyme's own Duplicated type.
Calling Duplicated on any Flux model which was defined using @layer will allocate space for the gradient,
and passing that to gradient (or withgradient, or train!) will then use Enzyme instead of Zygote.
The gradient functions still return the gradient as usual, which can then be passed to update!:

The gradient returned here is also stored within dup_model.
Both share the same arrays -- what is returned is not a copy, just a view of the same memory (wrapped in NamedTuples instead of structs).
They will all be set to zero when you call gradient again, then replaced with the new values.
Alternatively, gradient(f, args...; zero=false) will add the new gradient to what's already stored.
Writing Const(x1) is optional, just plain x1 is implicitly constant.
Any set of Duplicated and Const arguments may appear in any order, so long as there is at least one Duplicated.
The gradient grads_f[1] can be passed to update! as usual.
But for convenience, you may also use what is stored within Duplicated.
These are equivalent ways to perform an update step:

Instead of using these FLux functions, you can also use Enzyme's own functions directly.
Enzyme.gradient works like this:

Note that what Enzyme.gradient returns is an object like deepcopy(model) of the same type, grads_e[1] isa Chain.
But its fields contain the same gradient.

Enzyme.jl has its own extensive documentation.

Second-order AD  Copy Chapter
If you calculate a gradient within the loss function, then training will involve 2nd derivatives.
While this is in principle supported by Zygote.jl, there are many bugs, and Enzyme.jl is probably a better choice.
---
Src/Reference/Training/Optimisers  Copy Section

Optimisation Rules
Any optimization rule from Optimisers.jl can be used with train! and
other training functions.
For full details of how the interface works, see the Optimisers.jl documentation.

Optimisers Reference  Copy Chapter
All optimisers return an object that, when passed to train!, will update the parameters passed to it.

Composing Optimisers  Copy Chapter
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.

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

It is possible to compose optimisers for some added flexibility.

Decays  Copy Chapter
Similar to optimisers, Flux also defines some simple decays that can be used in conjunction with other optimisers, or standalone.

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

---
Src/Reference/Training/Reference  Copy Section

Training API Reference
The new version of Flux's training code was written as an independent package, Optimisers.jl.
Only the function train! belongs to Flux itself.
The Optimisers package is designed to allow for immutable objects. But at present all Flux models contain parameter arrays (such as Arrays and CuArrays) which can be updated in-place.
Because of this:
* The objects returned by Optimisers.update! can be ignored.
* Flux defines its own version of setup which checks this assumption.
  (Using instead Optimisers.setup will also work, they return the same thing.)
The available optimization rules are listed the optimisation rules page here. See the Optimisers documentation for details on how the rules work.

train! uses @progress which should show a progress bar in VSCode automatically.
To see one in a terminal, you will need to install TerminalLoggers.jl
and follow its setup instructions.

There is also a method of train! which similarly takes Duplicated(model) and uses Enzyme.jl for differentiation (see (@ref autodiff-enzyme)):

Optimisation Modifiers  Copy Chapter
The state returned by setup can be modified to temporarily prevent training of
some parts of the model, or to change the learning rate or other hyperparameter.
The functions for doing so may be accessed as Flux.freeze!, Flux.thaw!, and Flux.adjust!.
All mutate the state (or part of it) and return nothing.

---
Src/Reference/Destructure  Copy Section

Flat vs. Nested Structures

A Flux model is a nested structure, with parameters stored within many layers. Sometimes you may want a flat representation of them, to interact with functions expecting just one vector. This is provided by destructure:

Both destructure and the Restructure function can be used within gradient computations. For instance, this computes the Hessian ∂²L/∂θᵢ∂θⱼ of some loss function, with respect to all parameters of the Flux model. The resulting matrix has off-diagonal entries, which cannot really be expressed in a nested structure:

In order to collect all parameters of a model into a list instead, you can use the trainables function:

Any mutation of the elements of the resulting list will affect the model's parameters.
All Parameters  Copy Chapter
The functions destructure and trainables live in Optimisers.jl.