Flax NNX

https://flax.readthedocs.io/en/latest/
Admin
Flax NNX is a simplified API for Flax, a flexible end-to-end neural network library for JAX,...
Tokens:820,225
Snippets:5,245
Trust Score:8.3
Update:1 month ago
Show doc for...
Tokens:
Raw
### Setup imports and dummy data

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

Initial setup for Flax NNX and JAX 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 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]]
```

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

### Install Flax

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

Install the Flax library using pip. This is the initial setup step.

```bash
# ! pip install -U flax
```

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

### 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(...)

```

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

### Bridge Configuration Example

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

Basic configuration setup for the bridge component.

```python
bridge = Bridge(config="default")
bridge.initialize()
```

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

### 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)
```

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

### Initialize Bridge Component

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

Basic setup for the bridge component. Ensure all dependencies are installed before execution.

```python
from flax import bridge

# Initialize the bridge
bridge_instance = bridge.Bridge(config="default_config.yaml")
bridge_instance.start()
```

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

### Bridge Implementation Example

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

A basic example demonstrating how to set up a bridge component.

```python
import flax
# Bridge implementation code here
```

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

### 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
```

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

### MiniGPT Example

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

A basic example of using MiniGPT.

```python
import torch
from minigpt import MiniGPT

model = MiniGPT()
output = model.generate("Hello, world!")
print(output)
```

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

### MiniGPT Example

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

A basic example of using MiniGPT.

```python
import minigpt

model = minigpt.load_model()
output = model.generate("Hello, world!")
print(output)
```

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

### Install Flax

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

Install the Flax library in your environment.

```bash
!pip install -q flax
```

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

### MNIST Training Setup

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

Initial setup and training loop for the MNIST dataset using Flax.

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

class CNN(nn.Module):
  @nn.compact
  def __call__(self, x):
    x = nn.Conv(features=32, kernel_size=(3, 3))(x)
    x = nn.relu(x)
    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
    x = x.reshape((x.shape[0], -1))
    x = nn.Dense(features=10)(x)
    return x

def create_train_state(rng, learning_rate):
  cnn = CNN()
  params = cnn.init(rng, jnp.ones([1, 28, 28, 1]))['params']
  tx = optax.adam(learning_rate)
  return train_state.TrainState.create(apply_fn=cnn.apply, params=params, tx=tx)

@jax.jit
def train_step(state, batch):
  def loss_fn(params):
    logits = state.apply_fn({'params': params}, batch['image'])
    loss = optax.softmax_cross_entropy_with_integer_labels(logits=logits, labels=batch['label']).mean()
    return loss
  grad_fn = jax.value_and_grad(loss_fn)
  loss, grads = grad_fn(state.params)
  state = state.apply_gradients(grads=grads)
  return state, loss
```

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

### Install Flax from repository

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

Command to install the latest development version of Flax directly from GitHub.

```bash
pip install git+https://github.com/google/flax.git
```

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

### Install Flax and JAX

Source: https://flax.readthedocs.io/en/latest/mnist_tutorial.html

Use pip to install Flax and JAX with CUDA support if needed. Uncomment these lines when working in environments like Google Colab or Jupyter Notebook.

```bash
# !pip install -U "jax[cuda12]"
# !pip install -U flax
```

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

### Install Flax and dependencies with pip

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

Installs the latest Flax without dependencies, plus jaxtyping, kagglehub, and treescope. Run these commands in a Colab or Jupyter environment.

```python
! pip install --no-deps -U flax
! pip install jaxtyping kagglehub treescope
```

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

### InstanceNorm Example Usage

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

Demonstrates the usage of InstanceNorm and its equivalence to LayerNorm and GroupNorm. Requires Flax, JAX, and NumPy. The example initializes InstanceNorm and compares its output to LayerNorm and GroupNorm with equivalent configurations.

```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))
print(nnx.state(layer, nnx.Param))
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)
```

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

### Install required packages for image segmentation

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

Install OpenCV, Grain, Albumentations, Pillow, Flax, Optax, and Orbax checkpoint manager. These packages are needed for data loading, augmentation, and model training.

```bash
!pip install -U opencv-python-headless grain albumentations Pillow
!pip install -U flax optax orbax-checkpoint
```

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

### Image Segmentation with U-Net on Oxford Pets

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

This notebook demonstrates a complete image segmentation workflow using a simple U-Net on the Oxford Pets dataset. It covers data loading, model definition, training, and visualization of predictions. The final code block shows how to display a few validation examples with the helper function.

```python
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import torch
from torchvision import transforms
from torch.utils.data import Dataset, DataLoader
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision.models import resnet50, ResNet50_Weights
import os
import random
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
from tqdm import tqdm
import warnings
warnings.filterwarnings('ignore')

# Set device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')
```

```python
def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

set_seed()
```

```python
class OxfordPetDataset(Dataset):
    def __init__(self, root_dir, split='train', transform=None, target_transform=None):
        self.root_dir = root_dir
        self.split = split
        self.transform = transform
        self.target_transform = target_transform
        self.images_dir = os.path.join(root_dir, 'images')
        self.masks_dir = os.path.join(root_dir, 'annotations', 'trimaps')
        self.image_files = [f for f in os.listdir(self.images_dir) if f.endswith('.jpg')]
        self.image_files.sort()
        # Split the dataset
        train_files, temp_files = train_test_split(self.image_files, test_size=0.3, random_state=42)
        val_files, test_files = train_test_split(temp_files, test_size=0.5, random_state=42)
        if split == 'train':
            self.image_files = train_files
        elif split == 'val':
            self.image_files = val_files
        elif split == 'test':
            self.image_files = test_files

    def __len__(self):
        return len(self.image_files)

    def __getitem__(self, idx):
        img_name = self.image_files[idx]
        img_path = os.path.join(self.images_dir, img_name)
        mask_name = os.path.splitext(img_name)[0] + '.png'
        mask_path = os.path.join(self.masks_dir, mask_name)
        image = Image.open(img_path).convert('RGB')
        mask = Image.open(mask_path).convert('L')
        if self.transform:
            image = self.transform(image)
        if self.target_transform:
            mask = self.target_transform(mask)
        return image, mask
```

```python
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

target_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor()
])
```

```python
train_dataset = OxfordPetDataset(root_dir='data/oxford-pet', split='train', transform=transform, target_transform=target_transform)
val_dataset = OxfordPetDataset(root_dir='data/oxford-pet', split='val', transform=transform, target_transform=target_transform)
test_dataset = OxfordPetDataset(root_dir='data/oxford-pet', split='test', transform=transform, target_transform=target_transform)

batch_size = 16
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
```

```python
class UNet(nn.Module):
    def __init__(self, in_channels=3, out_channels=1):
        super(UNet, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(in_channels, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2)
        )
        self.middle = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2)
        )
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(64, out_channels, kernel_size=2, stride=2)
        )

    def forward(self, x):
        x1 = self.encoder(x)
        x2 = self.middle(x1)
        x3 = self.decoder(x2)
        return x3
```

```python
model = UNet().to(device)
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-4)

num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for images, masks in tqdm(train_loader, desc=f'Epoch {epoch+1}/{num_epochs}'):
        images = images.to(device)
        masks = masks.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, masks)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f'Epoch {epoch+1} loss: {running_loss/len(train_loader):.4f}')
```

```python
def display_image_mask_pred(image, mask, pred, label=''):
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
    axes[0].imshow(image.permute(1, 2, 0).cpu().numpy())
    axes[0].set_title('Image')
    axes[0].axis('off')
    axes[1].imshow(mask.squeeze().cpu().numpy(), cmap='gray')
    axes[1].set_title('Mask')
    axes[1].axis('off')
    axes[2].imshow(pred.squeeze().cpu().numpy(), cmap='gray')
    axes[2].set_title('Prediction')
    axes[2].axis('off')
    plt.suptitle(f'Image, Mask, Prediction{label}')
    plt.show()
```

```python
model.eval()
with torch.no_grad():
    images, masks = next(iter(val_loader))
    images = images.to(device)
    masks = masks.to(device)
    outputs = model(images)
    preds = torch.sigmoid(outputs)
    preds = (preds > 0.5).float()
```

```python
for img, mask, pred in zip(images[:4], masks[:4], preds[:4]):
    display_image_mask_pred(img, mask, pred, label=" (validation set)")
```

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

### Visualize ViT predictions on Food 101 test images

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

Visualizes predictions from a ViT model finetuned on a subset of the Food 101 dataset. Uses nnx.softmax to compute probabilities, argmax to get predicted labels, and matplotlib to display images with expected and predicted labels. Requires test_images, test_indices, expected_labels, preds, and the label mapping dictionaries.

```python
num_samples = len(test_indices)
names_map = train_dataset.features["label"].names

probas = nnx.softmax(preds, axis=1)
pred_labels = probas.argmax(axis=1)
fig, axs = plt.subplots(1, num_samples, figsize=(20, 10))
for i in range(num_samples):
    img, expected_label = test_images[i], expected_labels[i]
    pred_label = pred_labels[i].item()
    proba = probas[i, pred_label].item()
    if img.dtype in (np.float32, ):
        img = img.astype(np.float32)
    expected_label_str = names_map[inv_labels_mapping[expected_label]]
    pred_label_str = names_map[inv_labels_mapping[pred_label]]
    axs[i].set_title(f"Expected: {expected_label_str} vs \nPredicted: {pred_label_str}, P={proba:.2f}")
    axs[i].imshow(img)
```

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

### 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
```

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

### Initialize and run a module

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

Demonstrates how to instantiate the module and perform a forward pass.

```python
model = Linear(2, 3, rngs=nnx.Rngs(0))
y = model(jnp.ones((1, 2)))
print(y)
```

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

### Bridge Implementation Example

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

A basic implementation of the bridge component as described in the guide.

```python
class Bridge:
    def __init__(self, source, target):
        self.source = source
        self.target = target

    def connect(self):
        print(f"Connecting {self.source} to {self.target}")
```

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

### Initialize and run a module

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

Demonstrates how to instantiate the module and perform a forward pass.

```python
rngs = nnx.Rngs(0)
model = Linear(32, 64, rngs)
x = jnp.ones((1, 32))
y = model(x)
```

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

### Evaluate model on test data

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

Selects a set of test indices, loads the corresponding images and labels from the validation dataset, and runs the model to get predictions. Uses jax.set_mesh to shard inputs with FSDP and reshard the predictions.

```python
test_indices = [1, 250, 500, 750, 1000, 1234]

test_images = [val_dataset[i]["image"] for i in test_indices]
expected_labels = [val_dataset[i]["label"] for i in test_indices]

with jax.set_mesh(mesh):
    inputs = jnp.asarray(test_images, out_sharding=jax.P("fsdp"))
    preds = eval_model(inputs)
    preds = jax.sharding.reshard(preds, jax.P())
```

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

### 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}")
```

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

### Initialize JAX environment and devices

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

Imports necessary libraries and configures 8 CPU devices for testing if no GPU or TPU is available.

```python
from functools import partial

import jax
from jax import numpy as jnp
from jax.sharding import PartitionSpec as P, NamedSharding, AxisType
import optax
import flax
from flax import nnx

# Ignore this if you are already running on a TPU or GPU
if not jax._src.xla_bridge.backends_are_initialized():
  jax.config.update('jax_num_cpu_devices', 8)
print(f'You have 8 “fake” JAX devices now: {jax.devices()}')
```

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

### 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)
```

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

### Import Libraries and Set Up Environment

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

Imports necessary libraries including PyTorch, torchvision, and transformers. Sets up the device (GPU if available) and random seed for reproducibility.

```python
import torch
import torchvision
from torchvision import transforms
from transformers import ViTForImageClassification, ViTFeatureExtractor
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import requests
from torch.utils.data import DataLoader, Dataset
import os
import random

# Set device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')

# Set random seed for reproducibility
seed = 42
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(seed)
```

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

### 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)
```

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

### 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__)

```

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

### Initialize and run a module

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

Demonstrates how to initialize the module with random keys and perform a forward pass.

```python
model = Linear(10, 5, rngs=nnx.Rngs(0))
y = model(jnp.ones((1, 10)))
print(y.shape)
```

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

### Initialize and run a module

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

Demonstrates how to initialize the module with random keys and perform a forward pass.

```python
model = Linear(3, 2, rngs=nnx.Rngs(0))
y = model(jnp.ones((1, 3)))
print(y)
```

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

### 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)))

```

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

### Initialize and run a module

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

Demonstrates how to initialize the module with random keys and perform a forward pass.

```python
model = Linear(3, 2, rngs=nnx.Rngs(0))
x = jnp.ones((1, 3))
y = model(x)
print(y)
```

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

### 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,),)
  )
})

```

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

### 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)
```

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

### Naive partial initialization

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

This example shows how to initialize a model and swap in pre-trained parameters. Note that this approach may temporarily increase memory usage by creating parameters that are subsequently discarded.

```python
# Some pretrained model state
old_state = nnx.state(TwoLayerMLP(4, rngs=nnx.Rngs(0)))

simple_model = nnx.eval_shape(lambda: TwoLayerMLP(4, rngs=nnx.Rngs(42)))
print(f'Number of jax arrays in memory at start: {len(jax.live_arrays())}')
# In this line, extra kernel and bias is created inside the new LoRALinear!
```

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

### Install Flax via Pip

Source: https://flax.readthedocs.io/en/latest/index.html

Installs the Flax library using pip. This is the standard method for installing Flax in a Python environment.

```bash
pip install flax

```

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

### Launch TensorBoard

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

Starts the TensorBoard server to visualize training logs from the specified directory.

```python
%tensorboard --logdir $trace_dir --port 6006
```
Tokens:
Raw