### Clone and Install Equinox

Source: https://github.com/patrick-kidger/equinox/blob/main/CONTRIBUTING.md

Clone the repository, navigate to the directory, and install dependencies along with pre-commit hooks using `uv`.

```bash
git clone https://github.com/your-username-here/equinox.git
cd equinox
uv run prek install
```

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

### Import Libraries for BERT Example

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/bert.ipynb

Imports necessary libraries including Equinox, JAX, Optax, HuggingFace transformers and datasets, and others for the BERT language model example. Ensure these are installed before running.

```python
import functools
from collections.abc import Mapping

import einops  # https://github.com/arogozhnikov/einops
import equinox as eqx
import jax
import jax.numpy as jnp
import numpy as np
import optax  # https://github.com/deepmind/optax
from datasets import load_dataset  # https://github.com/huggingface/datasets
from jaxtyping import Array, Float, Int  # https://github.com/google/jaxtyping
from tqdm import notebook as tqdm  # https://github.com/tqdm/tqdm
from transformers import AutoTokenizer  # https://github.com/huggingface/transformers
```

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

### Example Usage of init-apply MLP

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/init_apply.ipynb

This example demonstrates how to use the `make_mlp` function to create an MLP, initialize its parameters, and apply it to sample data, showing that parameter updates affect the output.

```python
import jax.numpy as jnp
import jax.random as jrandom
import jax.tree_util as jtu


def main(in_size=2, seed=5678):
    key = jrandom.PRNGKey(seed)

    init_fn, apply_fn = make_mlp(
        in_size=in_size, out_size=1, width_size=8, depth=1, key=key
    )

    x = jnp.arange(in_size)  # sample data
    params = init_fn()
    y1 = apply_fn(params, x)
    params = jtu.tree_map(lambda p: p + 1, params)  # "stochastic gradient descent"
    y2 = apply_fn(params, x)
    assert y1 != y2


main()
```

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

### Install Equinox using pip

Source: https://github.com/patrick-kidger/equinox/blob/main/README.md

Install the Equinox library using pip. Requires Python 3.10+.

```bash
pip install equinox
```

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

### Vmap'd Stateful Layer Example

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/api/nn/stateful.md

Illustrates how to create and use vmap'd stateful layers by extending the `Counter` example. This involves managing state for batched operations. Requires `jax.random`.

```python
import jax.random as jr

class Model(eqx.Module):
    linear: eqx.nn.Linear
    counter: Counter
    v_counter: Counter

    def __init__(self, key):
        # Not-stateful layer
        self.linear = eqx.nn.Linear(2, 2, key=key)
        # Stateful layer.
        self.counter = Counter()
        # Vmap'd stateful layer. (Whose initial state will include a batch dimension.)
        self.v_counter = eqx.filter_vmap(Counter, axis_size=2)()

    def __call__(self, x: Array, state: eqx.nn.State) -> tuple[Array, eqx.nn.State]:
        # This bit happens as normal.
        assert x.shape == (2,)
        x = self.linear(x)
        x, state = self.counter(x, state)

        # For the vmap, we have to restrict our state to just those states we want to
        # vmap, and then update the overall state again afterwards.
        #
        # After all, the state for `self.counter` isn't expecting to be batched, so we
        # have to remove that.
        substate = state.substate(self.v_counter)
        x, substate = eqx.filter_vmap(self.v_counter)(x, substate)
        state = state.update(substate)

        return x, state

key = jr.PRNGKey(0)
model, state = eqx.nn.make_with_state(Model)(key)
x = jnp.array([5.0, -1.0])
model(x, state)
```

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

### Import Libraries

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/train_rnn.ipynb

Imports necessary libraries for Equinox, JAX, NumPy, and Optax. Ensure these libraries are installed.

```python
import math

import equinox as eqx
import jax
import jax.lax as lax
import jax.numpy as jnp
import jax.random as jrandom
import numpy as np
import optax  # https://github.com/deepmind/optax
```

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

### Calculate Example Loss and Inference Output

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/mnist.ipynb

Computes the loss for a given model and data, and performs inference to get predictions. Note that running JAX operations outside of a JIT'd region is slow and intended for exploration.

```python
# Example loss
loss_value = loss(model, dummy_x, dummy_y)
print(loss_value.shape)  # scalar loss
# Example inference
output = jax.vmap(model)(dummy_x)
print(output.shape)  # batch of predictions
```

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

### Setup for Parallel Computation

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/parallelism.ipynb

Imports necessary libraries and configures JAX to use 8 CPU devices for demonstration. Sets up hyperparameters, generates synthetic data, initializes a simple MLP model, and an Adam optimizer.

```python
import equinox as eqx
import jax
import jax.numpy as jnp
import jax.random as jr
import jax.sharding as jshard
import numpy as np
import optax  # https://github.com/deepmind/optax


# Force 8 CPUs for our toy problem
jax.config.update("jax_platform_name", "cpu")
jax.config.update("jax_num_cpu_devices", 8)

# Hyperparameters
dataset_size = 64
channel_size = 4
hidden_size = 32
depth = 1
learning_rate = 3e-4
num_steps = 100
batch_size = 16  # must be a multiple of our number of devices.

# Generate some synthetic data
xs = np.random.normal(size=(dataset_size, channel_size))
ys = np.sin(xs)

model = eqx.nn.MLP(channel_size, channel_size, hidden_size, depth, key=jr.PRNGKey(6789))
optim = optax.adam(learning_rate)
opt_state = optim.init(eqx.filter(model, eqx.is_inexact_array))
```

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

### Custom Stateful Layer Example

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/api/nn/stateful.md

Demonstrates how to create a custom stateful layer using `equinox.nn.StateIndex`. This layer increments a counter each time it's called. Requires `equinox` and `jax.numpy`.

```python
import equinox as eqx
import jax.numpy as jnp
from jaxtyping import Array

class Counter(eqx.Module):
    index: eqx.nn.StateIndex

    def __init__(self):
        init_state = jnp.array(0)
        self.index = eqx.nn.StateIndex(init_state)

    def __call__(self, x: Array, state: eqx.nn.State) -> tuple[Array, eqx.nn.State]:
        value = state.get(self.index)
        new_x = x + value
        new_state = state.set(self.index, value + 1)
        return new_x, new_state

counter, state = eqx.nn.make_with_state(Counter)()
x = jnp.array(2.3)

num_calls = state.get(counter.index)
print(f"Called {num_calls} times.")  # 0

_, state = counter(x, state)
num_calls = state.get(counter.index)
print(f"Called {num_calls} times.")  # 1

_, state = counter(x, state)
num_calls = state.get(counter.index)
print(f"Called {num_calls} times.")  # 2
```

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

### Use Equinox Model with JAX Transformations

Source: https://github.com/patrick-kidger/equinox/blob/main/README.md

Demonstrates how a model defined with Equinox is fully compatible with JAX operations like `jit` and `grad`. This example calculates gradients for a loss function.

```python
@jax.jit
@jax.grad
def loss_fn(model, x, y):
    pred_y = jax.vmap(model)(x)
    return jax.numpy.mean((y - pred_y) ** 2)

batch_size, in_size, out_size = 32, 2, 3
model = Linear(in_size, out_size, key=jax.random.PRNGKey(0))
x = jax.numpy.zeros((batch_size, in_size))
y = jax.numpy.zeros((batch_size, out_size))
grads = loss_fn(model, x, y)
```

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

### Define init-apply functions for Equinox MLP

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/init_apply.ipynb

This function partitions an Equinox MLP into its inexact parameters and static components, returning `init_fn` to get the parameters and `apply_fn` to perform the forward pass.

```python
import equinox as eqx


def make_mlp(in_size, out_size, width_size, depth, *, key):
    mlp = eqx.nn.MLP(
        in_size, out_size, width_size, depth, key=key
    )  # insert your model here
    params, static = eqx.partition(mlp, eqx.is_inexact_array)

    def init_fn():
        return params

    def apply_fn(_params, x):
        model = eqx.combine(_params, static)
        return model(x)

    return init_fn, apply_fn
```

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

### Import Libraries for ViT

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/vision_transformer.ipynb

Imports necessary libraries including Equinox, JAX, NumPy, Optax, and PyTorch for dataset loading and transformations. Ensure these libraries are installed before running.

```python
import functools

import einops  # https://github.com/arogozhnikov/einops
import equinox as eqx
import jax
import jax.numpy as jnp
import jax.random as jr
import numpy as np
import optax  # https://github.com/deepmind/optax

# We'll use PyTorch to load the dataset.
import torch
import torchvision
import torchvision.transforms as transforms
from jaxtyping import Array, Float, PRNGKeyArray
```

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

### RNN Training Output

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/train_rnn.ipynb

Example output logs from an RNN training session, showing step-by-step loss values. Monitor these values to assess model convergence.

```text
step=0, loss=0.6816999316215515
step=1, loss=0.7202574014663696
step=2, loss=0.6925007104873657
step=3, loss=0.689198911190033
step=4, loss=0.6808685064315796
step=5, loss=0.7059305906295776
step=6, loss=0.6922754049301147
step=7, loss=0.6842439770698547
step=8, loss=0.6972116231918335
step=9, loss=0.7047306299209595
step=10, loss=0.6993851661682129
step=11, loss=0.6921849846839905
step=12, loss=0.6844913959503174
step=13, loss=0.6941200494766235
step=14, loss=0.6870629787445068
step=15, loss=0.6922240257263184
step=16, loss=0.6966875195503235
step=17, loss=0.7021255493164062
step=18, loss=0.6913468241691589
step=19, loss=0.6915531158447266
step=20, loss=0.6906869411468506
step=21, loss=0.6945821046829224
step=22, loss=0.6963403820991516
step=23, loss=0.6893304586410522
step=24, loss=0.6923031210899353
step=25, loss=0.6952496767044067
step=26, loss=0.6937462687492371
step=27, loss=0.6946915984153748
step=28, loss=0.6912715435028076
step=29, loss=0.6945470571517944
step=30, loss=0.6928573250770569
step=31, loss=0.6918295621871948
step=32, loss=0.6926039457321167
step=33, loss=0.691811203956604
step=34, loss=0.696336567401886
step=35, loss=0.693527340888977
step=36, loss=0.6909832954406738
step=37, loss=0.6898350715637207
step=38, loss=0.693118691444397
step=39, loss=0.6962690353393555
step=40, loss=0.6943768262863159
step=41, loss=0.6929119229316711
step=42, loss=0.6921533942222595
step=43, loss=0.6970506906509399
step=44, loss=0.6914128065109253
step=45, loss=0.6925110220909119
step=46, loss=0.6876767873764038
step=47, loss=0.6977562308311462
step=48, loss=0.6887734532356262
step=49, loss=0.6956733465194702
step=50, loss=0.6988524198532104
step=51, loss=0.6972949504852295
step=52, loss=0.6935367584228516
step=53, loss=0.6899304389953613
step=54, loss=0.6940433979034424
step=55, loss=0.6932569742202759
step=56, loss=0.6964170932769775
step=57, loss=0.6952816843986511
step=58, loss=0.6925933361053467
step=59, loss=0.700016975402832
step=60, loss=0.6929588317871094
step=61, loss=0.6919406652450562
step=62, loss=0.6893216371536255
step=63, loss=0.6881398558616638
step=64, loss=0.6941375136375427
step=65, loss=0.6908596754074097
step=66, loss=0.6938614845275879
step=67, loss=0.6939255595207214
step=68, loss=0.691447377204895
step=69, loss=0.6932423114776611
step=70, loss=0.6937750577926636
step=71, loss=0.691257119178772
step=72, loss=0.6900532245635986
step=73, loss=0.6922309398651123
step=74, loss=0.6899502277374268
step=75, loss=0.6930654048919678
step=76, loss=0.6942011117935181
step=77, loss=0.6899413466453552
step=78, loss=0.6950610876083374
step=79, loss=0.6900242567062378
step=80, loss=0.691747784614563
step=81, loss=0.6899303793907166
step=82, loss=0.6910462379455566
step=83, loss=0.69475257396698
step=84, loss=0.6886341571807861
step=85, loss=0.6912660598754883
step=86, loss=0.6889529824256897
step=87, loss=0.6940121054649353
step=88, loss=0.6970347762107849
step=89, loss=0.687224268913269
step=90, loss=0.6900577545166016
step=91, loss=0.6913183927536011
step=92, loss=0.6916753649711609
step=93, loss=0.6899659633636475
step=94, loss=0.6911211013793945
step=95, loss=0.694290041923523
step=96, loss=0.7031664848327637
step=97, loss=0.6912339925765991
step=98, loss=0.6968348026275635
step=99, loss=0.6970176100730896
step=100, loss=0.6857004165649414
step=101, loss=0.6842451095581055
step=102, loss=0.6882964968681335
step=103, loss=0.6855384111404419
step=104, loss=0.6909692287445068
step=105, loss=0.6905874013900757
step=106, loss=0.6900045871734619
step=107, loss=0.6865564584732056
step=108, loss=0.6820229887962341
step=109, loss=0.6879786849021912
step=110, loss=0.6853011846542358
step=111, loss=0.68475741147995
step=112, loss=0.682267427444458
step=113, loss=0.6880433559417725
step=114, loss=0.6814002990722656
step=115, loss=0.6823583841323853
step=116, loss=0.6794727444648743
step=117, loss=0.6785068511962891
step=118, loss=0.6811013221740723
step=119, loss=0.6747442483901978
step=120, loss=0.6660218238830566
step=121, loss=0.6700407266616821
step=122, loss=0.6526561975479126
step=123, loss=0.6608943939208984
step=124, loss=0.6293025612831116
step=125, loss=0.6483496427536011
step=126, loss=0.6219364404678345
step=127, loss=0.5961954593658447
step=128, loss=0.6002600193023682
step=129, loss=0.5647848844528198
step=130, loss=0.5256890058517456
step=131, loss=0.510317325592041
step=132, loss=0.47984960675239563
step=133, loss=0.5084915161132812
step=134, loss=0.4301827549934387
step=135, loss=0.4290550649166107
step=136, loss=0.3755859136581421
step=137, loss=0.2937808036804199
step=138, loss=0.26023393869400024
step=139, loss=0.23048073053359985
step=140, loss=0.21439003944396973
step=141, loss=0.1652923822402954
step=142, loss=0.1283920854330063
step=143, loss=0.10732141137123108
step=144, loss=0.09533026814460754
step=145, loss=0.0801059827208519
step=146, loss=0.06879423558712006
step=147, loss=0.05884774401783943
```

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

### Python: Diamond Inheritance Pattern Example

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/pattern.md

Demonstrates a diamond inheritance pattern for building a differential equation solver using multiple abstract base classes. This structure adheres to Equinox's design principles.

```python
class AbstractSolver(eqx.Module):
    @abc.abstractmethod
    def step(...):
        raise NotImplementedError

class AbstractAdaptiveSolver(AbstractSolver):
    tolerance: eqx.AbstractVar[float]

class AbstractImplicitSolver(AbstractSolver):
    root_finder: eqx.AbstractVar[AbstractRootFinder]

class ImplicitEuler(AbstractAdaptiveSolver, AbstractImplicitSolver):
    tolerance: float
    root_finder: AbstractRootFinder = Newton()

    def step(...):
        ...

solver = ImplicitEuler(tolerance=1e-3)
```

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

### Load CIFAR10 Dataset with Torchvision

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/vision_transformer.ipynb

Loads the CIFAR10 dataset using torchvision, applying specified transformations for training and testing. Ensure torchvision and torch are installed.

```python
transform_train = transforms.Compose(
    [
        transforms.RandomCrop(32, padding=4),
        transforms.Resize((height, width)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ]
)

transform_test = transforms.Compose(
    [
        transforms.Resize((height, width)),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ]
)

train_dataset = torchvision.datasets.CIFAR10(
    "CIFAR",
    train=True,
    download=True,
    transform=transform_train,
)

test_dataset = torchvision.datasets.CIFAR10(
    "CIFAR",
    train=False,
    download=True,
    transform=transform_test,
)

trainloader = torch.utils.data.DataLoader(
    train_dataset, batch_size=batch_size, shuffle=True, drop_last=True
)

testloader = torch.utils.data.DataLoader(
    test_dataset, batch_size=batch_size, shuffle=True, drop_last=True
)
```

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

### Data Sharding and Model Parallelism Setup

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/parallelism.ipynb

Defines sharding configurations for data and model parameters using JAX's mesh and NamedSharding. This is essential for distributing computations across multiple devices.

```python
num_devices = len(jax.devices())
mesh = jax.make_mesh(
    (num_devices,), ("batch",), axis_types=(jax.sharding.AxisType.Auto,)
)
data_sharding = jshard.NamedSharding(mesh, jshard.PartitionSpec("batch"))
model_sharding = jshard.NamedSharding(mesh, jshard.PartitionSpec())
```

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

### Training Loop with Stateful Model

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/stateful.ipynb

Execute the training loop, repeatedly calling `make_step` to update the model, state, and optimizer state. This example uses full-batch gradient descent for simplicity.

```python
dataset_size = 10
learning_rate = 3e-4
steps = 5
seed = 5678

key = jr.PRNGKey(seed)
mkey, xkey, xkey2 = jr.split(key, 3)

model, state = eqx.nn.make_with_state(Model)(mkey)

xs = jr.normal(xkey, (dataset_size, 3))
ys = jnp.sin(xs) + 1
optim = optax.adam(learning_rate)
opt_state = optim.init(eqx.filter(model, eqx.is_inexact_array))

for _ in range(steps):
    # Full-batch gradient descent in this simple example.
    model, state, opt_state = make_step(model, state, opt_state, xs, ys)
```

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

### Define a model with stateful layers

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/stateful.ipynb

This is a placeholder comment indicating the start of a model definition that includes a mix of stateful and non-stateful layers. The actual model architecture would follow.

```python
# This model is just a weird mish-mash of stateful and non-stateful layers for
```

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

### Running the Training Loop with Data Sharding

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/parallelism.ipynb

Iterates through the training data, shards each batch, and performs a training step. Explicitly shards data outside the JIT region to guide the compiler.

```python
model, opt_state = eqx.filter_shard((model, opt_state), model_sharding)
for step, (x, y) in \
    zip(
        range(1, num_steps + 1),
        train_dataloader((xs, ys), batch_size),
    ):
    x, y = eqx.filter_shard((x, y), data_sharding)
    model, opt_state = train_step(model, opt_state, x, y)
```

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

### Python: Example of Non-Cooperative Multiple Inheritance Issue

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/pattern.md

Illustrates a common pitfall in Python's multiple inheritance where `super().__init__` is not called, leading to unexpected behavior. This highlights why Equinox avoids `super()`.

```python
class A:
    def __init__(self, x):
        self.x = x
        # Not calling super().__init__, because the superclass is just `object`, right?

class AA:
    def __init__((...):
        super().__init__(...)
        ...

class B(A, AA):
    pass

B()
```

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

### Serve Documentation Locally

Source: https://github.com/patrick-kidger/equinox/blob/main/CONTRIBUTING.md

Build and serve the documentation locally using `mkdocs serve` to preview changes. Access the local copy at `localhost:8000`.

```bash
uv run mkdocs serve
```

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

### Main Training Loop Configuration

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/score_based_diffusion.ipynb

Sets up the model, optimizer, and data loading, then iterates through the training steps. Includes hyperparameters for model architecture, optimization, and sampling.

```python
def main(
    # Model hyperparameters
    patch_size=4,
    hidden_size=64,
    mix_patch_size=512,
    mix_hidden_size=512,
    num_blocks=4,
    t1=10.0,
    # Optimisation hyperparameters
    num_steps=1_000_000,
    lr=3e-4,
    batch_size=256,
    print_every=10_000,
    # Sampling hyperparameters
    dt0=0.1,
    sample_size=10,
    # Seed
    seed=5678,
):
    key = jr.PRNGKey(seed)
    model_key, train_key, loader_key, sample_key = jr.split(key, 4)
    data = mnist()
    data_mean = jnp.mean(data)
    data_std = jnp.std(data)
    data_max = jnp.max(data)
    data_min = jnp.min(data)
    data_shape = data.shape[1:]
    data = (data - data_mean) / data_std

    model = Mixer2d(
        data_shape,
        patch_size,
        hidden_size,
        mix_patch_size,
        mix_hidden_size,
        num_blocks,
        t1,
        key=model_key,
    )
    int_beta = lambda t: t  # Try experimenting with other options here!
    weight = lambda t: 1 - jnp.exp(
        -int_beta(t)
    )  # Just chosen to upweight the region near t=0.

    opt = optax.adabelief(lr)
    # Optax will update the floating-point JAX arrays in the model.
    opt_state = opt.init(eqx.filter(model, eqx.is_inexact_array))

    total_value = 0
    total_size = 0
    for step, data in zip(
        range(num_steps),
        dataloader(data, batch_size, key=loader_key)
    ):
        value, model, train_key, opt_state = make_step(
            model, weight, int_beta, data, t1, train_key, opt_state, opt.update
        )
        total_value += value.item()
        total_size += 1
```

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

### Python: Concrete Class Example

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/pattern.md

An example of a concrete class that has no abstract methods or attributes, making it instantiable. This follows the 'concrete means final' principle.

```python
# This is clearly something we can instantiate: it has no abstract methods/attributes.
class ConcreteArray(eqx.Module):
    def some_method(self):
        pass

# This is clearly also without abstract methods/attributes, and also doesn't break the
```

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

### equinox.debug.get_num_traces

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/api/debug.md

Gets the current number of traces.

```APIDOC
## equinox.debug.get_num_traces

### Description
Gets the current number of traces.

### Parameters
(No parameters explicitly documented)

### Request Example
(No request example provided)

### Response
(No response details provided)
```

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

### Custom Weight Initialization for Linear Layer

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/tricks.md

Shows how to replace the weight of a specific linear layer with custom initialization using `equinox.tree_at`. This is useful for non-standard parameter initialization strategies.

```python
linear = eqx.nn.Linear(...)
new_weight = jax.random.normal(...)
where = lambda l: l.weight
new_linear = eqx.tree_at(where, linear, new_weight)
```

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

### Main Function Execution

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/score_based_diffusion.ipynb

This snippet simply calls the main function to initiate the process.

```python
main()
```

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

### Vision Transformer Initialization and Training Execution

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/vision_transformer.ipynb

Initializes the Vision Transformer model, an AdamW optimizer, and the optimizer state. It then calls the `train` function to start the training process.

```python
key = jr.PRNGKey(2003)

model = VisionTransformer(
    embedding_dim=embedding_dim,
    hidden_dim=hidden_dim,
    num_heads=num_heads,
    num_layers=num_layers,
    dropout_rate=dropout_rate,
    patch_size=patch_size,
    num_patches=num_patches,
    num_classes=num_classes,
    key=key,
)

optimizer = optax.adamw(
    learning_rate=lr,
    b1=beta1,
    b2=beta2,
)

state = optimizer.init(eqx.filter(model, eqx.is_inexact_array))

model, state, losses = train(model, optimizer, state, trainloader, num_steps, key=key)
```

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

### Import necessary libraries

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/stateful.ipynb

Imports the required libraries for building and training Equinox models, including JAX for numerical computation and Optax for optimization.

```python
import equinox as eqx
import jax
import jax.numpy as jnp
import jax.random as jr
import optax  # https://github.com/deepmind/optax
```

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

### Import necessary libraries for JAX and Equinox

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/mnist.ipynb

Imports core libraries including Equinox for model building, JAX for numerical computation, Optax for training, PyTorch for data handling, and jaxtyping for type annotations.

```python
import equinox as eqx
import jax
import jax.numpy as jnp
import optax  # https://github.com/deepmind/optax
import torch  # https://pytorch.org
import torchvision  # https://pytorch.org
from jaxtyping import Array, Float, Int, PyTree  # https://github.com/google/jaxtyping
```

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

### Initialize Model and State

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/stateful.ipynb

Use `eqx.nn.make_with_state` to initialize both the model parameters and the initial state. This function separates the model's parameters from its state, returning them as distinct PyTrees.

```python
model, state = eqx.nn.make_with_state(Model)(mkey)
```

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

### Prepare SST2 Dataset and Tokenizer

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/bert.ipynb

Sets up the tokenizer and prepares the SST2 dataset for training. It maps the tokenizer function to the dataset, ensuring sentences are tokenized and formatted correctly for JAX, including padding and truncation.

```python
tokenizer = AutoTokenizer.from_pretrained(
    "google/bert_uncased_L-2_H-128_A-2", model_max_length=128
)


def tokenize(example):
    return tokenizer(example["sentence"], padding="max_length", truncation=True)


ds = load_dataset("sst2")
ds = ds.map(tokenize, batched=True)
ds.set_format(type="jax", columns=["input_ids", "token_type_ids", "label"])
```

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

### Initialize AdamW Optimizer

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/mnist.ipynb

Initializes the AdamW optimizer from Optax with a specified learning rate. This is a common first step in setting up the training process.

```python
optim = optax.adamw(LEARNING_RATE)
```

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

### Import Libraries for Equinox

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/frozen_layer.ipynb

Imports essential libraries including Equinox, JAX, and Optax for numerical computation and optimization.

```python
import equinox as eqx
import jax
import jax.numpy as jnp
import jax.random as jrandom
import jax.tree_util as jtu
import optax  # https://github.com/deepmind/optax
```

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

### Concrete Class Implementing Intermediate ABC

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/pattern.md

Implement a concrete class that defines the abstract attributes declared in an intermediate ABC. This example shows how to initialize abstract attributes and call the superclass constructor.

```python
class AbstractPolynomialInterpolation(AbstractInterpolation)
    coeffs: Array

    def __init__(self, coeffs: Array):
        self.coeffs = coeffs

    def degree(self) -> int:
        return len(self.coeffs)

    def __call__(self, x: Array) -> Array:
        return jnp.polyval(self.coeffs, x)


class CubicInterpolation(AbstractPolynomialInterpolation):
    def __init__(self, ts: Array, xs: Array):
        coeffs = ...  # some implementation
        super().__init__(coeffs)
```

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

### JIT-compiled Training Step with Sharding and Memory Donation

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/parallelism.ipynb

A JIT-compiled training step that shards model and data, computes gradients, updates parameters, and donates memory for efficiency. Use `donate='all'` to allow JAX to reuse memory for inputs and outputs.

```python
@eqx.filter_jit(donate="all")
def train_step(model, opt_state, x, y):
    model, opt_state = eqx.filter_shard((model, opt_state), model_sharding)
    x, y = eqx.filter_shard((x, y), data_sharding)
    grads = eqx.filter_grad(compute_loss)(model, x, y)
    updates, opt_state = optim.update(grads, opt_state)
    model = eqx.apply_updates(model, updates)
    return eqx.filter_shard((model, opt_state), model_sharding)
```

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

### Partitioning Model Parameters

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/mnist.ipynb

Use `eqx.partition` to separate model components into parameters and static parts. This is necessary when differentiating functions that include non-trainable elements.

```python
# This will work!
params, static = eqx.partition(model, eqx.is_array)

def loss2(params, static, x, y):
    model = eqx.combine(params, static)
    return loss(model, x, y)


loss_value, grads = jax.value_and_grad(loss2)(params, static, dummy_x, dummy_y)
print(loss_value)
```

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

### Define Training and Evaluation Steps

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/bert.ipynb

Defines the core functions for training and evaluation using JAX and Equinox. `compute_loss` calculates the loss and gradients, while `make_step` applies these gradients for model updates. `make_eval_step` is used for evaluating the model's performance.

```python
@eqx.filter_value_and_grad
def compute_loss(classifier, inputs, key):
    batch_size = inputs["token_ids"].shape[0]
    batched_keys = jax.random.split(key, num=batch_size)
    logits = jax.vmap(classifier, in_axes=(0, None, 0))(inputs, True, batched_keys)
    return jnp.mean(
        optax.softmax_cross_entropy_with_integer_labels(
            logits=logits, labels=inputs["label"]
        )
    )


def make_step(model, inputs, opt_state, key, tx):
    key, new_key = jax.random.split(key)
    loss, grads = compute_loss(model, inputs, key)
    grads = jax.lax.pmean(grads, axis_name="devices")

    updates, opt_state = tx.update(grads, opt_state, model)
    model = eqx.apply_updates(model, updates)
    return loss, model, opt_state, new_key


def make_eval_step(model, inputs):
    return jax.vmap(functools.partial(model, enable_dropout=False))(inputs)


p_make_eval_step = eqx.filter_pmap(make_eval_step)
```

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

### Gradient Descent Update for Neural Network

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/all-of-equinox.md

Performs a gradient descent update on the model parameters using a specified learning rate. This example demonstrates how to apply computed gradients to update the model, leveraging JAX's tree manipulation utilities.

```python
x_key, y_key, model_key = jax.random.split(jax.random.PRNGKey(0), 3)
# Example data
x = jax.random.normal(x_key, (100, 2))
y = jax.random.normal(y_key, (100, 2))
model = NeuralNetwork(model_key)
# Compute gradients
grads = loss(model, x, y)
# Perform gradient descent
learning_rate = 0.1
new_model = jax.tree_util.tree_map(lambda m, g: m - learning_rate * g, model, grads)
```

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

### Training Step with Stateful Model

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/stateful.ipynb

Implement a training step using `jax.jit` and `equinox.filter_grad`. This function computes gradients, updates model parameters, and threads the state and optimizer state through the training loop.

```python
@eqx.filter_jit
def make_step(model, state, opt_state, xs, ys):
    grads, state = eqx.filter_grad(compute_loss, has_aux=True)(model, state, xs, ys)
    updates, opt_state = optim.update(grads, opt_state)
    model = eqx.apply_updates(model, updates)
    return model, state, opt_state
```

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

### Define GAN Hyperparameters

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/deep_convolutional_gan.ipynb

Sets up key hyperparameters for training the GAN, including learning rate, batch size, training steps, image dimensions, and latent space size.

```python
# Hyperparameters
lr = 0.0002
beta1 = 0.5
beta2 = 0.999
batch_size = 32
num_steps = 100000
image_size = (64, 64, 1)
height, width, channels = image_size
latent_size = 100
```

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

### Import Libraries for GAN

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/deep_convolutional_gan.ipynb

Imports necessary libraries for building and training a GAN in Equinox, including JAX, Optax, and PyTorch for data loading.

```python
from collections.abc import Callable

import equinox as eqx
import jax
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
import optax

# We'll use PyTorch to load the dataset.
import torch
import torchvision
import torchvision.transforms as transforms
```

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

### Prepare for Inference with Stateful Model

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/stateful.ipynb

Use `eqx.nn.inference_mode` to set stateful layers (like BatchNorm) to inference behavior. Wrap the model and state using `eqx.Partial` to create an inference-ready callable.

```python
inference_model = eqx.nn.inference_mode(model)
inference_model = eqx.Partial(inference_model, state=state)
```

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

### Configure Optimizer and Parallel Training

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/bert.ipynb

Configures the Adam optimizer with gradient clipping and initializes the optimizer state. It also sets up `pmap` for parallel training across multiple devices, ensuring the batch size is compatible with the number of available devices.

```python
epochs = 3
batch_size = 32
learning_rate = 1e-5

num_devices = jax.device_count()
assert batch_size % num_devices == 0, (
    "The batch size must be a multiple of the number of devices."
)

tx = optax.adam(learning_rate=learning_rate)
tx = optax.chain(optax.clip_by_global_norm(1.0), tx)
opt_state = tx.init(classifier_chkpt)

p_make_step = eqx.filter_pmap(functools.partial(make_step, tx=tx), axis_name="devices")

```

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

### Partition and Combine Model with JAX Transformations

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/all-of-equinox.md

Demonstrates using `eqx.partition` and `eqx.combine` to separate JAX arrays from other PyTree leaves for use with JAX transformations like `jax.jit` and `jax.grad`. The `eqx.is_array` filter function is used to partition the model.

```python
@ft.partial(jax.jit, static_argnums=1)  # `static` must be a PyTree of non-arrays.
@jax.grad  # differentiates with respect to `params`, as it is the first argument
def loss(params, static, x, y):
    model = eqx.combine(params, static)
    pred_y = jax.vmap(model)(x)
    return jax.numpy.mean((y - pred_y) ** 2)

params, static = eqx.partition(model, eqx.is_array)
loss(params, static, x, y)
```

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

### Replicate Model and Data Across Devices

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/bert.ipynb

Replicates the optimizer state, model checkpoint, and training data across all available local devices. This is a prerequisite for distributed training.

```python
opt_state = jax.device_put_replicated(opt_state, jax.local_devices())
model = jax.device_put_replicated(classifier_chkpt, jax.local_devices())
train_key = jax.device_put_replicated(train_key, jax.local_devices())
```

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

### Simple Dataloader

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/parallelism.ipynb

Provides a basic dataloader implementation that randomly slices the dataset and shuffles between epochs. This implementation uses NumPy for efficiency with small datasets that fit in host memory.

```python
# Simple dataloader; randomly slices our dataset and shuffles between epochs.
# In NumPy for speed, as our dataset is small enough to fit entirely in host memory.
#
# For larger datasets (that require loading from disk) then use PyTorch's `DataLoader`
```

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

### equinox.nn.Conv

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/api/nn/conv.md

Documentation for the general Conv layer, including its constructor and call method.

```APIDOC
## equinox.nn.Conv

### Description
Represents a general convolutional layer.

### Methods
- `__init__`: Initializes the Conv layer.
- `__call__`: Applies the convolution operation.
```

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

### GAN Model Initialization and Training Call

Source: https://github.com/patrick-kidger/equinox/blob/main/docs/examples/deep_convolutional_gan.ipynb

Initializes the generator and discriminator models, their states, and optimizers. It then calls the `train` function to train the GAN. Ensure `latent_size`, `image_size`, `lr`, `beta1`, `beta2`, `dataloader`, `num_steps`, and `jr` (JAX random) are defined.

```python
key = jr.PRNGKey(2003)

key, gen_key, disc_key = jr.split(key, 3)

generator = Generator(input_shape=latent_size, output_shape=image_size, key=gen_key)
discriminator = Discriminator(input_shape=image_size, key=disc_key)

generator_state = eqx.nn.State(generator)
discriminator_state = eqx.nn.State(discriminator)

generator_optimizer = optax.adam(lr, b1=beta1, b2=beta2)
discriminator_optimizer = optax.adam(lr, b1=beta1, b2=beta2)

generator_opt_state = generator_optimizer.init(eqx.filter(generator, eqx.is_array))
discriminator_opt_state = discriminator_optimizer.init(
    eqx.filter(discriminator, eqx.is_array)
)

(
    generator,
    discriminator,
    generator_state,
    discriminator_state,
    generator_losses,
    discriminator_losses,
    key,
) = train(
    generator,
    discriminator,
    generator_optimizer,
    discriminator_optimizer,
    generator_opt_state,
    discriminator_opt_state,
    generator_state,
    discriminator_state,
    dataloader,
    num_steps,
    key,
)
```