### Setup dummy variables for checkpointing example

Source: https://flax.readthedocs.io/en/latest/guides/converting_and_upgrading/orbax_upgrade_guide.html

Defines constants and sample data structures for demonstrating checkpoint migration.

```python
# Create some dummy variables for this example.
MAX_STEPS = 5
CKPT_PYTREE = [12, {'bar': np.array((2, 3))}, [1, 4, 10]]
TARGET_PYTREE = [0, {'bar': np.array((0))}, [0, 0, 0]]
```

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

### Module Setup Example

Source: https://flax.readthedocs.io/en/latest/developer_notes/module_lifecycle.html

Illustrates how submodules are typically initialized within the setup method of a Flax Module. This method is called on a bound Module, providing access to the full Module API.

```python
class AutoEncoder(nn.Module):
  def setup(self):
    self.encoder = Encoder(...)
    self.decoder = Decoder(...)

```

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

### Setup Imports and Dummy Data

Source: https://flax.readthedocs.io/en/latest/_sources/guides/jax_and_nnx_transforms.rst

Initializes necessary imports from Flax and JAX, and generates random input data for subsequent examples.

```python
from flax import nnx
import jax

x = jax.random.normal(jax.random.key(0), (1, 2))
y = jax.random.normal(jax.random.key(1), (1, 3))
```

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

### Setup: Import dependencies and define model

Source: https://flax.readthedocs.io/en/latest/_sources/guides/checkpointing.ipynb

Imports necessary libraries, sets up a checkpoint directory, and defines an example `TwoLayerMLP` model by subclassing `nnx.Module`. This is a prerequisite for checkpointing operations.

```python
from flax import nnx
import orbax.checkpoint as ocp
import jax
from jax import numpy as jnp
import numpy as np

ckpt_dir = ocp.test_utils.erase_and_create_empty('/tmp/my-checkpoints/')
```

```python
class TwoLayerMLP(nnx.Module):
  def __init__(self, dim, rngs: nnx.Rngs):
    self.linear1 = nnx.Linear(dim, dim, rngs=rngs, use_bias=False)
    self.linear2 = nnx.Linear(dim, dim, rngs=rngs, use_bias=False)

  def __call__(self, x):
    x = self.linear1(x)
    return self.linear2(x)

# Instantiate the model and show we can run it.
model = TwoLayerMLP(4, rngs=nnx.Rngs(0))
x = jax.random.normal(jax.random.key(42), (3, 4))
assert model(x).shape == (3, 4)
```

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

### Setup Code for Flax Examples

Source: https://flax.readthedocs.io/en/latest/flip/1009-optimizer-api.html

Provides necessary imports and utility functions for running Flax code snippets, including model definition, data preprocessing, loss calculation, and learning rate scheduling.

```python
import functools
from typing import Callable, Sequence

import jax
import jax.numpy as jnp
import flax
import flax.linen as nn
import tensorflow as tf
import tensorflow_datasets as tfds


def pp(features):
  return {
      'image': tf.cast(features['image'], tf.float32) / 255 - 0.5,
      'label': features['label'],
  }


class Model(nn.Module):

  @nn.compact
  def __call__(self, inputs):
    x = inputs.reshape([inputs.shape[0], -1])
    x = nn.normalization.BatchNorm(True)(x)
    x = nn.Dense(10)(x)
    x = nn.log_softmax(x)
    return x


def onehot(labels, num_classes, on_value=1.0, off_value=0.0):
  x = (labels[..., None] == jnp.arange(num_classes)[None])
  x = jax.lax.select(
      x, jnp.full(x.shape, on_value), jnp.full(x.shape, off_value))
  return x.astype(jnp.float32)

def xent_loss(logits, labels):
  return -jnp.sum(
      onehot(labels, num_classes=10) * logits) / labels.size

def get_learning_rate(step):
  return 0.1


model = Model()
rng = jax.random.key(0)
ds = tfds.load('mnist')['train'].take(160).map(pp).batch(16)
batch = next(iter(ds))
variables = model.init(rng, jnp.array(batch['image'][:1]))
jax.tree_util.tree_map(jnp.shape, variables)
```

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

### Using Flax's `setup` for Initialization Logic

Source: https://flax.readthedocs.io/en/latest/_sources/examples/gemma.ipynb

Demonstrates the use of the `setup` method in Flax modules for defining submodules or parameters during initialization. This is an alternative to using `@nn.compact`.

```python
class SetupModel(nn.Module):
    features: Sequence[int]

    def setup(self):
        # Define submodules in setup
        self.layers = [
            nn.Dense(feat) for feat in self.features
        ]

    def __call__(self, x):
        for i, layer in enumerate(self.layers):
            x = layer(x)
            if i != len(self.features) - 1:
                x = nn.relu(x)
        return x
```

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

### InstanceNorm Example Usage

Source: https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/normalization.html

Demonstrates how to use InstanceNorm and compares its output with LayerNorm and GroupNorm for identical inputs and configurations. Ensure you have JAX and NumPy installed.

```python
>>> from flax import nnx
>>> import jax
>>> import numpy as np
>>> # dimensions: (batch, height, width, channel)
>>> x = jax.random.normal(jax.random.key(0), (2, 3, 4, 5))
>>> layer = nnx.InstanceNorm(5, rngs=nnx.Rngs(0))
>>> nnx.state(layer, nnx.Param)
State({
  'bias': Param( # 5 (20 B)
    value=Array([0., 0., 0., 0., 0.], dtype=float32)
  ),
  'scale': Param( # 5 (20 B)
    value=Array([1., 1., 1., 1., 1.], dtype=float32)
  )
})
>>> y = layer(x)
>>> # having a channel_axis of -1 in InstanceNorm is identical to reducing all non-batch,
>>> # non-channel axes and using the feature_axes as the feature_axes in LayerNorm
>>> y2 = nnx.LayerNorm(5, reduction_axes=[1, 2], feature_axes=-1, rngs=nnx.Rngs(0))(x)
>>> np.testing.assert_allclose(y, y2, atol=1e-7)
>>> y3 = nnx.GroupNorm(5, num_groups=x.shape[-1], rngs=nnx.Rngs(0))(x)
>>> np.testing.assert_allclose(y, y3, atol=1e-7)
```

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

### Module Setup Lifecycle

Source: https://flax.readthedocs.io/en/latest/developer_notes/module_lifecycle.html

Demonstrates that `setup` is not called when a top-level Module is constructed, and attributes assigned in `setup` are not available until the module is bound.

```python
class TopLevelAccess(nn.Module):

  def setup(self):
    self.foo = nn.Dense(2)

mdl = TopLevelAccess()
assert not hasattr(mdl, "foo")  # foo is not defined because setup is not called
```

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

### Set up virtual environment

Source: https://flax.readthedocs.io/en/latest/_sources/contributing.md

Use this command to create a virtual environment for installing dependencies. Activate it using `. env/bin/activate`.

```bash
python3 -m virtualenv env
. env/bin/activate
```

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

### Basic Image Segmentation Model

Source: https://flax.readthedocs.io/en/latest/_sources/examples/image_segmentation.ipynb

A foundational example of an image segmentation model. Ensure you have the necessary libraries installed.

```python
import flax.linen as nn
import jax
import jax.numpy as jnp

class SegmentationModel(nn.Module):
    num_classes: int

    @nn.compact
    def __call__(self, x, training: bool = False):
        # Example encoder layers
        x = nn.Conv(features=32, kernel_size=(3, 3))(x)
        x = nn.relu(x)
        x = nn.max_pool(x, ksize=(2, 2), strides=(2, 2))

        # Example decoder layers
        x = nn.ConvTranspose(features=32, kernel_size=(3, 3), strides=(2, 2))(x)
        x = nn.relu(x)

        # Output layer
        logits = nn.Conv(features=self.num_classes, kernel_size=(1, 1))(x)
        return logits

# Example usage:
key = jax.random.PRNGKey(0)
model = SegmentationModel(num_classes=10)
input_shape = (1, 256, 256, 3)  # Batch, Height, Width, Channels
params = model.init(key, jnp.zeros(input_shape))['params']

# Dummy input
input_data = jnp.zeros(input_shape)
output = model.apply({'params': params}, input_data)
print(output.shape)  # Expected output shape: (1, 256, 256, 10)
```

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

### Training and Evaluation Loop Setup

Source: https://flax.readthedocs.io/en/latest/examples/vit_training.html

Sets up the main training and evaluation loops using `tqdm` for progress bars. Includes logic for data loading, device placement, and metric updates.

```python
import tqdm


bar_format = "{desc}[{n_fmt}/{total_fmt}]{postfix} [{elapsed}<{remaining}]"

# We define a view of the model sharing the weights but with attributes set for evaluation
eval_model = nnx.view(model, deterministic=True)
rngs = nnx.Rngs(12)

def train_one_epoch(epoch):
    with tqdm.tqdm(
        desc=f"[train] epoch: {epoch + 1}/{num_epochs}, ",
        total=total_steps,
        bar_format=bar_format,
        leave=True,
    ) as pbar, jax.set_mesh(mesh):
        prev_loss = None
        for batch in train_loader:

            # Convert np.ndarray to jax.Array on GPUs
            images = jax.device_put(batch["image"].astype(act_dtype), device=jax.P("fsdp"))
            labels = jax.device_put(batch["label"].astype(int), device=jax.P("fsdp"))

            loss = train_step(model, optimizer, rngs, (images, labels))
            if prev_loss is not None:
                # Async metrics recording and printing
                train_metrics_history["train_loss"].append(prev_loss.item())
                pbar.set_postfix({"loss": prev_loss.item()})
            prev_loss = loss
            pbar.update(1)


def evaluate_model(epoch):
    # Computes the metrics on the training and test sets after each training epoch.
    with jax.set_mesh(mesh):
        eval_metrics.reset()  # Reset the eval metrics
        for val_batch in val_loader:

            # Convert np.ndarray to jax.Array on GPUs
            images = jax.device_put(val_batch["image"].astype(act_dtype), device=jax.P("fsdp"))
            labels = jax.device_put(val_batch["label"].astype(int), device=jax.P("fsdp"))

            eval_step(eval_model, (images, labels), eval_metrics)

    for metric, value in eval_metrics.compute().items():
        eval_metrics_history[f'val_{metric}'].append(value)

    print(f"[val] epoch: {epoch + 1}/{num_epochs}")
    print(f"- total loss: {eval_metrics_history['val_loss'][-1]:0.4f}")
    print(f"- Accuracy: {eval_metrics_history['val_accuracy'][-1]:0.4f}")
```

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

### Basic Image Segmentation Model

Source: https://flax.readthedocs.io/en/latest/_sources/examples/image_segmentation.ipynb

A foundational example of an image segmentation model in Flax. This snippet demonstrates the core architecture and setup.

```python
import flax.linen as nn
import jax
import jax.numpy as jnp

class SegmentationModel(nn.Module):
    num_classes: int

    @nn.compact
    def __call__(self, x, training: bool = False):
        # Example encoder-decoder structure (simplified)
        # Convolutional layers for feature extraction
        x = nn.Conv(features=32, kernel_size=(3, 3), padding='SAME')(x)
        x = nn.relu(x)
        x = nn.max_pool(x, ksize=(2, 2), strides=(2, 2))

        # ... more encoder layers ...

        # Example decoder layers
        # Transposed convolutions or upsampling for segmentation map generation
        x = nn.ConvTranspose(features=32, kernel_size=(3, 3), strides=(2, 2), padding='SAME')(x)
        x = nn.relu(x)

        # ... more decoder layers ...

        # Final convolutional layer to predict class scores
        logits = nn.Conv(features=self.num_classes, kernel_size=(1, 1))(x)
        return logits

# Example usage:
key = jax.random.PRNGKey(0)
model = SegmentationModel(num_classes=10)

# Dummy input data (batch_size, height, width, channels)
input_data = jnp.ones((1, 256, 256, 3))

# Initialize model parameters
params = model.init(key, input_data)['params']

# Perform a forward pass
output = model.apply({'params': params}, input_data)

print(f"Output shape: {output.shape}")
```

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

### Basic Image Segmentation Model

Source: https://flax.readthedocs.io/en/latest/_sources/examples/image_segmentation.ipynb

A foundational example of an image segmentation model in Flax. This snippet demonstrates the core architecture and setup.

```python
import flax.linen as nn
import jax.numpy as jnp

class SegmentationModel(nn.Module):
    num_classes: int

    @nn.compact
    def __call__(self, x, training: bool = False):
        # Example encoder layers (simplified)
        x = nn.Conv(features=32, kernel_size=(3, 3), padding='SAME')(x)
        x = nn.relu(x)
        x = nn.max_pool(x, ksize=(2, 2), strides=(2, 2))

        # Example decoder layers (simplified)
        x = nn.ConvTranspose(features=32, kernel_size=(3, 3), strides=(2, 2), padding='SAME')(x)
        x = nn.relu(x)
        x = nn.Conv(features=self.num_classes, kernel_size=(1, 1))(x)
        return x

```

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

### Main Function

Source: https://flax.readthedocs.io/en/latest/_sources/mnist_tutorial.ipynb

Sets up the training process, including model initialization, optimizer configuration, and starting the training loop.

```python
def main(argv):
  if len(argv) > 1:  # pytype skips the specific flags
    raise app.UsageError('Too many command-line arguments.')

  # Create input iterators
  train_ds, eval_ds, num_train_examples, num_eval_examples = create_input_iterators()

  # Initialize model and training state
  key = jax.random.PRNGKey(0)
  model = CNN(num_classes=10)
  dummy_input = jnp.zeros((1, 28, 28, 1))
  params = model.init(key, dummy_input)['params']
  tx = optax.adam(learning_rate=0.001)
  state = create_train_state(model.apply, params, tx)

  # Train the model
  train(state, train_ds, eval_ds, num_epochs=10)

if __name__ == '__main__':
  app.run(main)
```

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

### Implement a Training Loop

Source: https://flax.readthedocs.io/en/latest/_sources/hijax/hijax.ipynb

Sets up a model and optimizer, then defines and executes a JIT-compiled training step. This example shows how to handle model parameters and gradients for optimization.

```python
# hijax Variables by default
model = Block(2, 64, 3, rngs=nnx.Rngs(0))
optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)

@jax.jit
def train_step(model, optimizer, x, y):
  graphdef, params, nondiff = nnx.split(model, nnx.Param, ...)
  def loss_fn(params):
    model =  nnx.merge(graphdef, params, nondiff)
    return ((model(x) - y) ** 2).mean()

  loss, grads = jax.value_and_grad(loss_fn)(nnx.with_vars(params, mutable=False))  # immutable for jax.grad
  optimizer.update(model, grads)

  return loss

for _ in range(3):
  loss = train_step(model, optimizer, x=jnp.ones((10, 2)), y=jnp.ones((10, 3)))
  print(f"{loss = !s}")
```

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

### Basic Neural Network Example

Source: https://flax.readthedocs.io/en/latest/_sources/examples/gemma.ipynb

A simple feedforward neural network implemented using Flax. Requires JAX and Flax to be installed.

```python
import jax
import jax.numpy as jnp
from flax import linen as nn

class MLP(nn.Module):
    features: list[int]

    @nn.compact
    def __call__(self, x):
        for i, feat in enumerate(self.features):
            x = nn.Dense(feat)(x)
            if i != len(self.features) - 1:
                x = nn.relu(x)
        return x

key = jax.random.PRNGKey(0)
model = MLP(features=[128, 64, 10])
variables = model.init(key, jnp.ones((1, 784)))
params = variables['params']

output = model.apply({'params': params}, jnp.ones((1, 784)))
print(output.shape)
```

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

### Initialize and Use Optimizer

Source: https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/training/optimizer.html

Demonstrates how to initialize an Optimizer with a model and an Optax transformation, and then use it to update model parameters based on gradients.

```python
>>> import jax, jax.numpy as jnp
>>> from flax import nnx
>>> import optax
...
>>> class Model(nnx.Module):
...   def __init__(self, rngs):
...     self.linear1 = nnx.Linear(2, 3, rngs=rngs)
...     self.linear2 = nnx.Linear(3, 4, rngs=rngs)
...   def __call__(self, x):
...     return self.linear2(self.linear1(x))
...
>>> x = jax.random.normal(jax.random.key(0), (1, 2))
>>> y = jnp.ones((1, 4))
...
>>> model = Model(nnx.Rngs(0))
>>> tx = optax.adam(1e-3)
>>> optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)
...
>>> loss_fn = lambda model: ((model(x) - y) ** 2).mean()
>>> loss_fn(model)
Array(2.3359997, dtype=float32)
>>> grads = nnx.grad(loss_fn)(model)
>>> _ = optimizer.update(model, grads)
>>> loss_fn(model)
Array(2.310461, dtype=float32)
```

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

### Initialize Model and Optimizer

Source: https://flax.readthedocs.io/en/latest/_sources/mnist_tutorial.ipynb

Initializes the CNN model, PRNG key, and the Adam optimizer. This setup is required before starting the training process.

```python
key = jax.random.PRNGKey(0)

model = CNN(num_classes=10)
params = model.init(key, jnp.zeros([1, 28, 28, 1]))['params']
opt = optim.Adam(learning_rate=0.001)
state = opt.create(params)

```

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

### Example Usage and Output

Source: https://flax.readthedocs.io/en/latest/guides/extracting_intermediates.html

Instantiates the `Foo` model and input data, then runs the `train_step` to demonstrate capturing intermediate values and gradients. The output shows the shapes of the captured states.

```python
model, x = Foo(4), jnp.ones((3, 4))
_, intermediates = train_step(model, x)
jax.tree.map(lambda a: a.shape, intermediates)

```

```text
State({
  'grad_of_x': Perturbation(
    value=(3, 4)
  ),
  'y': Intermediate(
    value=((3,),)
  )
})

```

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

### Haiku Forward Pass and Model Initialization

Source: https://flax.readthedocs.io/en/latest/migrating/haiku_to_flax.html

Defines the forward pass for the MLP and initializes the Haiku model. This setup is used for the scan-over-layers example.

```python
def forward(x, training: bool):
  return MLP(64, num_layers=5)(x, training)
model = hk.transform(forward)

sample_x = jnp.ones((1, 64))
params = model.init(jax.random.key(0), sample_x, training=False)
```

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

### Basic Image Segmentation with U-Net

Source: https://flax.readthedocs.io/en/latest/_sources/examples/image_segmentation.ipynb

Demonstrates a basic image segmentation setup using the U-Net architecture. Ensure you have the necessary libraries installed.

```python
import flax.linen as nn
import jax
import jax.numpy as jnp

class ConvBlock(nn.Module):
    features: int

    @nn.compact
    def __call__(self, x):
        x = nn.Conv(features=self.features, kernel_size=(3, 3), padding=[(1, 1), (1, 1)])(x)
        x = nn.relu(x)
        x = nn.Conv(features=self.features, kernel_size=(3, 3), padding=[(1, 1), (1, 1)])(x)
        x = nn.relu(x)
        return x


class UNet(nn.Module):
    features: list[int]

    @nn.compact
    def __call__(self, x):
        skip_connections = []
        for i, feat in enumerate(self.features):
            x = ConvBlock(features=feat)(x)
            if i < len(self.features) - 1:
                skip_connections.append(x)
                x = nn.max_pool(x, ksize=(2, 2), strides=(2, 2))

        for i, feat in enumerate(reversed(self.features[:-1])):
            x = nn.ConvTranspose2d(features=feat, kernel_size=(2, 2), strides=(2, 2))(x)
            x = jnp.concatenate([x, skip_connections.pop(-1)], axis=3)
            x = ConvBlock(features=feat)(x)

        x = nn.Conv(features=1, kernel_size=(1, 1))(x)
        x = nn.sigmoid(x)
        return x

```

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

### Basic JAX Training Loop Setup

Source: https://flax.readthedocs.io/en/latest/_sources/guides/data_loaders.ipynb

Sets up a basic JAX training loop structure with type annotations for images and labels. Requires importing JAX, NumPy, Flax, and jaxtyping.

```python
import jax
import jax.numpy as jnp
import numpy as np
from jaxtyping import Float, Int, Array
from flax import nnx

batch_size = 32

def train(model: nnx.Module, images:  Float[Array, "batch channels height width"], labels: Int[Array, "batch"]):
  pass

def train_loop(model):
  for images, labels in get_batches(train_ds):
    train(model, images, labels)
```

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

### Image Segmentation with JAX and Optax

Source: https://flax.readthedocs.io/en/latest/_sources/examples/image_segmentation.ipynb

Illustrates image segmentation using JAX for computation and Optax for optimization. This example requires Optax to be installed.

```python
import flax.linen as nn
import jax
import jax.numpy as jnp
import optax

class SimpleUNet(nn.Module):
    num_classes: int

    @nn.compact
    def __call__(self, x):
        # Downsampling path
        conv1 = nn.Conv(features=64, kernel_size=(3, 3), padding='same')(x)
        conv1 = nn.relu(conv1)
        pool1 = nn.max_pool(conv1, ksize=(2, 2), strides=(2, 2))

        conv2 = nn.Conv(features=128, kernel_size=(3, 3), padding='same')(pool1)
        conv2 = nn.relu(conv2)
        pool2 = nn.max_pool(conv2, ksize=(2, 2), strides=(2, 2))

        # Bottleneck
        bottleneck = nn.Conv(features=256, kernel_size=(3, 3), padding='same')(pool2)
        bottleneck = nn.relu(bottleneck)

        # Upsampling path (simplified)
        upsample1 = nn.ConvTranspose(features=128, kernel_size=(2, 2), strides=(2, 2))(bottleneck)
        concat1 = jnp.concatenate([upsample1, conv2], axis=-1) # Simplified concatenation
        conv3 = nn.Conv(features=128, kernel_size=(3, 3), padding='same')(concat1)
        conv3 = nn.relu(conv3)

        upsample2 = nn.ConvTranspose(features=64, kernel_size=(2, 2), strides=(2, 2))(conv3)
        concat2 = jnp.concatenate([upsample2, conv1], axis=-1) # Simplified concatenation
        conv4 = nn.Conv(features=64, kernel_size=(3, 3), padding='same')(concat2)
        conv4 = nn.relu(conv4)

        # Output layer
        output = nn.Conv(features=self.num_classes, kernel_size=(1, 1))(conv4)
        return output

key = jax.random.PRNGKey(0)

# Model initialization
model = SimpleUNet(num_classes=2) # Example: binary segmentation
input_shape = (1, 256, 256, 3) # (batch, height, width, channels)

params = model.init(key, jnp.zeros(input_shape))['params']

# Optimizer setup
optimizer = optax.adam(learning_rate=1e-3)
opt_state = optimizer.init(params)

# Dummy training step function (for demonstration)
def train_step(params, opt_state, batch):
    def loss_fn(params):
        predictions = model.apply({'params': params}, batch['image'])
        # Example loss: cross-entropy (requires one-hot encoded labels)
        labels_one_hot = jax.nn.one_hot(batch['label'], num_classes=model.num_classes)
        loss = optax.softmax_cross_entropy(predictions, labels_one_hot).mean()
        return loss

    grad_fn = jax.grad(loss_fn)
    grads = grad_fn(params)
    updates, opt_state = optimizer.update(grads, opt_state, params)
    params = optax.apply_updates(params, updates)
    return params, opt_state

# Dummy data
dummy_batch = {
    'image': jnp.ones(input_shape),
    'label': jnp.zeros((1, 256, 256, 1), dtype=jnp.int32)
}

# Perform one training step
new_params, new_opt_state = train_step(params, opt_state, dummy_batch)

print("Training step completed.")
```

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

### Basic Image Segmentation Model

Source: https://flax.readthedocs.io/en/latest/_sources/examples/image_segmentation.ipynb

A foundational example of an image segmentation model. This snippet demonstrates the core architecture and setup required for segmentation tasks.

```python
import flax.linen as nn
import jax
import jax.numpy as jnp

class ConvBlock(nn.Module):
    features: int

    @nn.compact
    def __call__(self, x):
        x = nn.Conv(self.features, (3, 3), padding=[(1, 1), (1, 1)])(x)
        x = nn.relu(x)
        x = nn.Conv(self.features, (3, 3), padding=[(1, 1), (1, 1)])(x)
        x = nn.relu(x)
        return x


class DownBlock(nn.Module):
    features: int

    @nn.compact
    def __call__(self, x):
        x = nn.max_pool(x, (2, 2), strides=(2, 2))
        return ConvBlock(self.features)(x)


class UpBlock(nn.Module):
    features: int

    @nn.compact
    def __call__(self, x, skip_connection):
        x = nn.ConvTranspose(self.features, (3, 3), strides=(2, 2), padding='same')(x)
        x = jnp.concatenate([x, skip_connection], axis=-1)
        return ConvBlock(self.features)(x)


class UNet(nn.Module):
    num_classes: int

    @nn.compact
    def __call__(self, x):
        skip_connections = []
        x = ConvBlock(32)(x)
        skip_connections.append(x)
        x = DownBlock(64)(x)
        skip_connections.append(x)
        x = DownBlock(128)(x)
        skip_connections.append(x)
        x = DownBlock(256)(x)
        skip_connections.append(x)
        x = DownBlock(512)(x)
        x = UpBlock(256)(x, skip_connections.pop())
        x = UpBlock(128)(x, skip_connections.pop())
        x = UpBlock(64)(x, skip_connections.pop())
        x = UpBlock(32)(x, skip_connections.pop())
        x = nn.Conv(self.num_classes, (1, 1))(x)
        return x

```

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

### Main execution block

Source: https://flax.readthedocs.io/en/latest/_sources/mnist_tutorial.ipynb

Set up the random number generator, create datasets, initialize the training state, and run the training loop.

```python
def main():
    # Create datasets
    ds_train, ds_test = create_datasets()

    # Initialize random number generator
    key = jax.random.PRNGKey(0)
    rng, init_rng = jax.random.split(key)

    # Initialize model and optimizer
    learning_rate = 0.001
    num_classes = 10
    state = create_train_state(init_rng, learning_rate, num_classes)

    # Train the model
    num_epochs = 5
    state = train_model(state, ds_train, ds_test, num_epochs)

if __name__ == '__main__':
    main()

```

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

### Import and Print Library Versions

Source: https://flax.readthedocs.io/en/latest/examples/image_segmentation.html

Imports core JAX, Flax, Optax, and Orbax libraries and prints their installed versions. This is useful for verifying the environment setup.

```python
import jax
import flax
import optax
import orbax.checkpoint as ocp
print("Jax version:", jax.__version__)
print("Flax version:", flax.__version__)
print("Optax version:", optax.__version__)
print("Orbax version:", ocp.__version__)

```

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

### Setup Toy Model and Data

Source: https://flax.readthedocs.io/en/latest/guides/optimization_cookbook.html

Initializes necessary JAX and Flax modules, defines a toy sequential model, a loss function, and generates fake data for training and evaluation.

```python
import jax
from flax import nnx
jax.config.update('jax_num_cpu_devices', 8)
import jax.numpy as jnp
import functools as ft
import optax
import matplotlib.pyplot as plt

lecun_normal = jax.nn.initializers.lecun_normal()

rngs = nnx.Rngs(0)

def make_model(rngs, init=lecun_normal):
    return nnx.Sequential(
        nnx.Linear(2,8, rngs=rngs, kernel_init=init),
        nnx.Linear(8,8, rngs=rngs, kernel_init=init))

def loss_fn(model, x, y):
    return jnp.mean((model(x) - y) ** 2)

```

```python
x = rngs.normal((32, 2))
y = rngs.normal((32, 8))

```

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

### Initialize and Use the U-Net Model

Source: https://flax.readthedocs.io/en/latest/_sources/examples/image_segmentation.ipynb

Demonstrates how to initialize the U-Net model parameters and perform a forward pass with dummy data. Requires JAX and Flax to be installed.

```python
import jax
import jax.random as random

key = random.PRNGKey(0)

# Assume input image shape is (batch_size, height, width, channels)
input_shape = (1, 256, 256, 3)

# Instantiate the model
model = UNet(num_classes=10)

# Initialize parameters
params = model.init(key, jnp.ones(input_shape))['params']

# Perform a forward pass
output = model.apply({'params': params}, jnp.ones(input_shape))

print(f"Output shape: {output.shape}")
# Expected output shape: (1, 256, 256, 10)
```

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

### Initialize a CNN Model

Source: https://flax.readthedocs.io/en/latest/_sources/examples/gemma.ipynb

Initializes a Convolutional Neural Network (CNN) model with specified input shape and random key. This is a common setup for starting training.

```python
key = jax.random.PRNGKey(0)
model = models.ResNet50(num_classes=10)
variables = model.init(key, jnp.zeros([1, 224, 224, 3]))["params"]
```

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

### Instantiate MLP Model and Prepare Inputs

Source: https://flax.readthedocs.io/en/latest/guides/parallel_training/flax_on_pjit.html

Sets up MLP hyperparameters, creates dummy input data, initializes a PRNG key, and defines an Optax optimizer. This prepares the environment for model training.

```python
# MLP hyperparameters.
BATCH, LAYERS, DEPTH, USE_SCAN = 8, 4, 1024, False
# Create fake inputs.
x = jnp.ones((BATCH, DEPTH))
# Initialize a PRNG key.
k = random.key(0)

# Create an Optax optimizer.
optimizer = optax.adam(learning_rate=0.001)
# Instantiate the model.
model = MLP(LAYERS, DEPTH, USE_SCAN)
```

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

### Get Hijax State Function

Source: https://flax.readthedocs.io/en/latest/_modules/flax/nnx/variablelib.html

Retrieves the state of a HijaxVariable, handling cases with and without QDD (Quantized Differential Data). Requires appropriate imports and setup for JAX.

```python
def _get_hijax_state(hijax_var: HijaxVariable | AbstractVariable) -> Variable:
  if hijax_var.has_qdd:
    tys: VariableQDD = jax.experimental.cur_qdd(hijax_var)
    leaf_vals = get_variable_p.bind(
      hijax_var,
      treedef=tys.treedef,
      avals=tuple(tys.leaf_avals),
      var_type=hijax_var._var_type,
      has_qdd=hijax_var.has_qdd,
    )
    variable = jax.tree.unflatten(tys.treedef, leaf_vals)
  else:
    assert hijax_var._treedef is not None
    assert hijax_var._leaves is not None
    if isinstance(hijax_var, (jax.core.Tracer, AbstractVariable)):
      leaf_avals = hijax_var._leaves
    else:
      leaf_avals = tuple(map(jax.typeof, hijax_var._leaves))
    leaf_vals = get_variable_p.bind(
      hijax_var,
      treedef=hijax_var._treedef,
      avals=leaf_avals,
      var_type=hijax_var._var_type,
      has_qdd=hijax_var.has_qdd,
    )
    variable = jax.tree.unflatten(hijax_var._treedef, leaf_vals)

  return variable
```

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

### Basic Image Segmentation Model

Source: https://flax.readthedocs.io/en/latest/_sources/examples/image_segmentation.ipynb

A foundational example demonstrating a simple image segmentation model architecture in Flax. This can be used as a starting point for custom segmentation tasks.

```python
import flax.linen as nn
import jax.numpy as jnp

class SegmentationModel(nn.Module):
    num_classes: int

    @nn.compact
    def __call__(self, x, training: bool = False):
        # Example encoder-decoder structure
        # Convolutional layers for feature extraction
        x = nn.Conv(features=32, kernel_size=(3, 3))(x)
        x = nn.relu(x)
        x = nn.max_pool(x, ksize=(2, 2), strides=(2, 2))

        # Add more layers as needed (e.g., more conv, pooling, upsampling)

        # Final convolutional layer to predict class scores
        logits = nn.Conv(features=self.num_classes, kernel_size=(1, 1))(x)
        return logits

```

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

### Initialize and Use the U-Net Model

Source: https://flax.readthedocs.io/en/latest/_sources/examples/image_segmentation.ipynb

Demonstrates how to initialize the U-Net model with dummy data and perform a forward pass. Ensure you have JAX and Flax installed.

```python
import jax
import jax.random as random

key = random.PRNGKey(0)

# Dummy input data (batch_size, height, width, channels)
dummy_input = jnp.ones((1, 256, 256, 3))

# Instantiate the model
model = UNet(num_classes=10) # Example: 10 classes

# Initialize model parameters
params = model.init(key, dummy_input)['params']

# Perform a forward pass
output = model.apply({'params': params}, dummy_input)

print(f"Output shape: {output.shape}")
```

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

### Initialize and Use the U-Net Model

Source: https://flax.readthedocs.io/en/latest/_sources/examples/image_segmentation.ipynb

Demonstrates how to initialize the U-Net model with dummy data and perform a forward pass. Ensure you have JAX and Flax installed.

```python
import jax
import jax.numpy as jnp
from flax.training import train_state
import optax

key = jax.random.PRNGKey(0)

# Dummy input data (batch_size, height, width, channels)
dummy_input = jnp.ones([1, 256, 256, 3])

# Initialize the model
model = UNet(num_classes=10)
params = model.init(key, dummy_input)['params']

# Create a dummy training state (optional, for training context)
optimizer = optax.adam(learning_rate=0.001)
training_state = train_state.TrainState.create(
    apply_fn=model.apply,
    params=params,
    tx=optimizer
)

# Perform a forward pass
output = model.apply({'params': training_state.params}, dummy_input, training=False)

print("Output shape:", output.shape)

```

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

### Initialize U-Net Model and Parameters

Source: https://flax.readthedocs.io/en/latest/_sources/examples/image_segmentation.ipynb

Demonstrates how to initialize the U-Net model and its parameters using a dummy input. This is a necessary step before training or inference.

```python
import jax

key = jax.random.PRNGKey(0)
model = UNet(num_classes=10)
# Dummy input shape (batch_size, height, width, channels)
dummy_input = jnp.zeros((1, 256, 256, 3))
params = model.init(key, dummy_input)['params']
print(f'Initialized model with {len(jax.tree_util.tree_leaves(params))} parameters.')

```

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

### Define AutoEncoder with Flax Linen

Source: https://flax.readthedocs.io/en/latest/migrating/linen_to_nnx.html

Define an auto-encoder model using Flax Linen, where layers are initialized lazily. This example shows the setup for encoder and decoder layers, and the encode, decode, and __call__ methods.

```python
class AutoEncoder(nn.Module):
  embed_dim: int
  output_dim: int

  def setup(self):
    self.encoder = nn.Dense(self.embed_dim)
    self.decoder = nn.Dense(self.output_dim)

  def encode(self, x):
    return self.encoder(x)

  def decode(self, x):
    return self.decoder(x)

  def __call__(self, x):
    x = self.encode(x)
    x = self.decode(x)
    return x

model = AutoEncoder(256, 784)
variables = model.init(jax.random.key(0), x=jnp.ones((1, 784)))

```

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

### Flax and JAX Setup Code

Source: https://flax.readthedocs.io/en/latest/_sources/flip/1009-optimizer-api.md

This code sets up a basic Flax model, loads MNIST data using TensorFlow Datasets, and initializes model variables. It's intended for running examples within this FLIP.

```python
import functools
from typing import Callable, Sequence

import jax
import jax.numpy as jnp
import flax
import flax.linen as nn
import tensorflow as tf
import tensorflow_datasets as tfds


def pp(features):
  return {
      'image': tf.cast(features['image'], tf.float32) / 255 - 0.5,
      'label': features['label'],
  }


class Model(nn.Module):

  @nn.compact
  def __call__(self, inputs):
    x = inputs.reshape([inputs.shape[0], -1])
    x = nn.normalization.BatchNorm(True)(x)
    x = nn.Dense(10)(x)
    x = nn.log_softmax(x)
    return x


def onehot(labels, num_classes, on_value=1.0, off_value=0.0):
  x = (labels[..., None] == jnp.arange(num_classes)[None])
  x = jax.lax.select(
      x, jnp.full(x.shape, on_value), jnp.full(x.shape, off_value))
  return x.astype(jnp.float32)


def xent_loss(logits, labels):
  return -jnp.sum(
      onehot(labels, num_classes=10) * logits) / labels.size


def get_learning_rate(step):
  return 0.1


model = Model()
rng = jax.random.key(0)
ds = tfds.load('mnist')['train'].take(160).map(pp).batch(16)
batch = next(iter(ds))
variables = model.init(rng, jnp.array(batch['image'][:1]))
jax.tree_util.tree_map(jnp.shape, variables)
```

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

### LinearGeneral Module Initialization Examples

Source: https://flax.readthedocs.io/en/latest/_modules/flax/nnx/nn/linear.html

Demonstrates various ways to initialize the LinearGeneral module with different input/output features and axes.

```python
layer = nnx.LinearGeneral(2, 4, rngs=nnx.Rngs(0))
```

```python
layer = nnx.LinearGeneral(2, (4, 5), rngs=nnx.Rngs(0))
```

```python
layer = nnx.LinearGeneral((2, 3), (4, 5), axis=(1, -1), rngs=nnx.Rngs(0))
```

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

### Extracting States from a Module in Flax NNX

Source: https://flax.readthedocs.io/en/latest/_modules/flax/nnx/graphlib.html

Demonstrates how to extract different types of states (parameters, batch statistics) from an NNX Module using the `nnx.state` function. It shows examples of getting a single state type, multiple state types separately, and all states.

```python
>>> from flax import nnx

>>> class Model(nnx.Module):
...   def __init__(self, rngs):
...     self.batch_norm = nnx.BatchNorm(2, rngs=rngs)
...     self.linear = nnx.Linear(2, 3, rngs=rngs)
...   def __call__(self, x):
...     return self.linear(self.batch_norm(x))

>>> model = Model(rngs=nnx.Rngs(0))
>>> # get the learnable parameters from the batch norm and linear layer
>>> params = nnx.state(model, nnx.Param)
>>> # get the batch statistics from the batch norm layer
>>> batch_stats = nnx.state(model, nnx.BatchStat)
>>> # get them separately
>>> params, batch_stats = nnx.state(model, nnx.Param, nnx.BatchStat)
>>> # get them together
>>> state = nnx.state(model)
```

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

### Initialize and Use the U-Net Model

Source: https://flax.readthedocs.io/en/latest/_sources/examples/image_segmentation.ipynb

Demonstrates how to initialize the U-Net model with dummy data and perform a forward pass. Ensure you have appropriate input dimensions and number of classes.

```python
import jax
import jax.random as random

key = random.PRNGKey(0)

# Example usage:
num_classes = 10
model = UNet(num_classes)

# Dummy input data (batch_size, height, width, channels)
# For example, a 256x256 image with 3 channels
input_shape = (1, 256, 256, 3)
dummy_input = jnp.zeros(input_shape, dtype=jnp.float32)

# Initialize model parameters
params = model.init(key, dummy_input)['params']

# Perform a forward pass
output = model.apply({'params': params}, dummy_input)

print(f"Output shape: {output.shape}") # Expected: (1, 256, 256, num_classes)
```

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

### Import dependencies and setup checkpoint directory

Source: https://flax.readthedocs.io/en/latest/guides/checkpointing.html

Import necessary libraries and set up a checkpoint directory. This is a prerequisite for saving and loading checkpoints.

```python
from flax import nnx
import orbax.checkpoint as ocp
import jax
from jax import numpy as jnp
import numpy as np

ckpt_dir = ocp.test_utils.erase_and_create_empty('/tmp/my-checkpoints/')

```

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

### Example Generated Text Output

Source: https://flax.readthedocs.io/en/latest/examples/minigpt.html

An example of the text output generated by the model.

```text
Once upon a timefect666PA carrier Louisiana denying►Files Mist tooltipbroken explo({666 Mens reprFiles Amit Sons explouriesodo explocmd termination �{\isans veins repr Tup Into terminationaudisansisans Cabinetdrivenbrokencmd tortured({isans Pall presum veins McCabe({({ Marlins terminationmananger dialogcmd termination explo brink ending SorcererHam reprications scattered Anna emblemisans infiltr Iranian Parkinsonelectric prin Sorcerergravityotor CLASS Kus({ Carm crest ExitHash terminationcmduriescmd 408 CLASSseys total technologiesoleon({ giftednette prinJacksonо�({ Anna674 408 "-brokenо� 36 DeVos Hound ratt targeted McCabe ample scattered gloveelectric Accord targeted Cabinet materially Option Kosovo blururies targetedHashcakes 408 pringravity pressedisans Claraidayssettinguries Wagner Rand explo 'entialbroken infiltr lottery Shatteredisansisans({ Untgravity62 terminationcmd Houndisans CLASSisans princesHash prin784 Thronesgravity targeted Parkinsonbroken Annafect grabs Addiction Carm targeted lottery CISisans CLASS gifted gifted Coldgravity StudentsTemporousUncommon Thrones incentivesisansisansisans({broken Carm Randisansisanselectric prinJacksonnox giftedications theirFilesExploregravity targeted Thrones Carm MG ParkinsonSurecakes incentivesmanncakes Yin sponsoredaurusudicrousの� agendaisansbasselectric desolategravity targeted aversionisans criminalmbudsman Thrones Carmisans Grassley technologies collectedisans threatened RISWithNoracticalNetwork Cabinet Balance Culture MAPHave wins subduedisans ClaraWithNo Clara Judges Judges Judges'
```

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

### Initialize and Use the U-Net Model

Source: https://flax.readthedocs.io/en/latest/_sources/examples/image_segmentation.ipynb

Demonstrates how to initialize the U-Net model with random parameters and perform a forward pass. Requires JAX and Flax to be installed.

```python
import jax
import jax.random as random

key = random.PRNGKey(0)

# Example input shape (batch, height, width, channels)
input_shape = (1, 256, 256, 3)

# Initialize the model
model = UNet(num_classes=10)
params = model.init(key, jnp.ones(input_shape), training=False)['params']

# Perform a forward pass
output = model.apply({'params': params}, jnp.ones(input_shape), training=False)

print(f'Output shape: {output.shape}')

```

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

### Initialize and Use the U-Net Model

Source: https://flax.readthedocs.io/en/latest/_sources/examples/image_segmentation.ipynb

Demonstrates how to initialize the U-Net model with random parameters and perform a forward pass. Requires JAX and Flax to be installed.

```python
import jax
import jax.random as random

key = random.PRNGKey(0)

# Assume input image shape is (batch_size, height, width, channels)
input_shape = (1, 256, 256, 3)

# Instantiate the model
model = UNet(num_classes=10) # Example with 10 classes

# Initialize parameters
params = model.init(key, jnp.ones(input_shape))['params']

# Perform a forward pass
output = model.apply({'params': params}, jnp.ones(input_shape))

print(f'Output shape: {output.shape}')
```

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

### Install Flax via pip

Source: https://flax.readthedocs.io/en/latest/_sources/index.rst

Standard installation command for Flax using pip.

```bash
pip install flax
```