Nabla

https://www.nablaml.com/
Admin
Nabla is a high-performance scientific computing and machine learning framework combining imperative...
Tokens:103,208
Snippets:1,085
Trust Score:5
Update:3 months ago
Show doc for...
Tokens:
Raw
### Setup Imports and Environment

Source: https://www.nablaml.com/_sources/examples/11_lora_finetuning.ipynb

Initial imports and setup for the Nabla LoRA/QLoRA fine-tuning example.

```python
from __future__ import annotations

import time
import shutil
from pathlib import Path

import numpy as np

import nabla as nb

print("Nabla LoRA/QLoRA Fine-Tuning example")
```

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

### Setup Development Environment

Source: https://www.nablaml.com/_sources/get_started.md

Clone the repository and install development dependencies for testing and linting.

```bash
git clone https://github.com/nabla-ml/nabla.git
cd nabla
python -m venv venv
source venv/bin/activate
pip install -r requirements-dev.txt
pip install -e ".[dev]"
```

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

### Importing Dependencies

Source: https://www.nablaml.com/_sources/examples/13_cnn_training.ipynb

Initial setup for the Nabla environment.

```python
from __future__ import annotations

import time

import numpy as np

import nabla as nb

print("Nabla CNN — Eager vs. Compiled")
```

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

### Initialize Imports

Source: https://www.nablaml.com/_sources/examples/04b_mlp_training_jax.ipynb

Setup necessary libraries and Nabla functional modules.

```python
import time
import numpy as np

import nabla as nb
from nabla.nn.functional import xavier_normal

print("Nabla MLP Training — JAX-style (functional)")
```

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

### Importing Nabla

Source: https://www.nablaml.com/_sources/examples/01_tensors_and_ops.ipynb

Initializes the library and verifies the installation.

```python
import numpy as np

import nabla as nb

print("Nabla imported successfully!")
```

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

### Nabla Output Example

Source: https://www.nablaml.com/examples/02_autodiff.html

Expected console output after running the initialization code.

```text
Nabla autodiff example
```

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

### Initialize Pipeline Configuration

Source: https://www.nablaml.com/_sources/examples/09_mlp_pipeline_inference.ipynb

Sets up the environment, imports, and constants for the pipeline inference example.

```python
import time
import numpy as np

import nabla as nb
from nabla import ops
from nabla.core.sharding import DeviceMesh, DimSpec, PartitionSpec as P
from nabla.ops import communication
from nabla.transforms import vmap

# Pipeline configuration
STAGES = 4
MICRO_BATCHES = 8
MICRO_BATCH_SIZE = 4
DIM = 16

print("Nabla Pipeline Inference example")
```

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

### Install Nabla via pip

Source: https://www.nablaml.com/_sources/get_started.md

Install the library using the Modular nightly index. Requires a virtual environment.

```bash
python -m venv venv
source venv/bin/activate
pip install --pre --extra-index-url https://whl.modular.com/nightly/simple/ modular nabla-ml
```

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

### Benchmark Output

Source: https://www.nablaml.com/examples/05_transforms_and_compile.html

Example console output showing performance metrics and compilation statistics.

```text
Compiled training loop:
Step     Loss
----------------------
10       0.937205
20       0.886352
30       0.833265
40       0.794533
50       0.760996
60       0.724098

Compiled cache stats: CompilationStats(hits=60, misses=1, fallbacks=0, hit_rate=98.4%)
Compiled avg step:   0.919 ms/step
Eager avg step:      21.697 ms/step
Compiled speedup:    23.6x
```

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

### Initialize Nabla Autodiff

Source: https://www.nablaml.com/_sources/examples/02_autodiff.ipynb

Basic setup for using Nabla's autodiff features.

```python
import numpy as np

import nabla as nb

print("Nabla autodiff example")
```

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

### Example Output of Compiled Training

Source: https://www.nablaml.com/examples/06b_transformer_jax

Displays the console output showing training progress and performance metrics.

```text
Compiled training:
Step     Loss
----------------------
10       1.1312
20       1.0764
30       1.0369
40       1.0024
50       0.9391
60       0.8277

Compiled cache stats: CompilationStats(hits=60, misses=1, fallbacks=0, hit_rate=98.4%)
Compiled avg step:   13.10 ms
Eager avg step:      209.53 ms
Compiled speedup:    16.0x
```

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

### JAX Benchmark Output

Source: https://www.nablaml.com/examples/10_jax_comparison_compiled.html

Example console output showing the results of the JAX JIT training benchmark.

```text
JAX warmup loss: 0.364887
  Step  50: loss = 0.127728

JAX JIT: 0.0463s (1294.7 steps/sec)
Loss: 0.337955 → 0.126996
```

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

### Performance Summary Output

Source: https://www.nablaml.com/examples/10_jax_comparison_compiled.html

Example output showing the performance metrics for Nabla and JAX execution modes.

```text
============================================================
PERFORMANCE SUMMARY
============================================================
Nabla @nb.compile:  0.0610s  (984.3 steps/sec)
Nabla Eager:        1.3166s  (45.6 steps/sec)
Nabla Eager (MAX):  4.7352s  (12.7 steps/sec)
JAX @jit:           0.0463s  (1294.7 steps/sec)

JAX JIT is 1.32x faster than Nabla compiled

Compile speedup over eager: 21.60x
Loss match (compiled vs eager): ✅ diff=0.00000000
```

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

### Implement an MLP using Module

Source: https://www.nablaml.com/api/nn/modules/index.html

Example of creating a multi-layer perceptron by inheriting from nabla.nn.Module.

```python
class MLP(nabla.nn.Module):
    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.fc1 = nabla.nn.Linear(in_dim, 64)
        self.fc2 = nabla.nn.Linear(64, out_dim)

    def forward(self, x):
        return self.fc2(nabla.relu(self.fc1(x)))

```

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

### Gradient Execution Output

Source: https://www.nablaml.com/examples/07_mlp_pipeline_parallel.html

Example output showing the timing and shape of the computed gradients.

```text
Gradient run time: 11.5356 s
Gradient avg run: 192.260 ms/run
Input gradient shape:  (8, 4, 16)
Weight gradient shape: (4, 16, 16)
Bias gradient shape:   (4, 16)
Weight grad range:     [-294.8232, 1527.5933]
```

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

### Verification Output

Source: https://www.nablaml.com/examples/07_mlp_pipeline_parallel.html

Example output confirming that the Nabla pipeline gradients match the JAX reference.

```text
Max input grad diff:  0.000061
Max weight grad diff: 0.000122
Max bias grad diff:   0.000015
✅ Nabla pipeline gradients match JAX sequential reference
```

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

### Configure 2D Sharding

Source: https://www.nablaml.com/examples/08_mlp_pp_dp_training.html

Setup sharding specifications for weights, inputs, and pipeline states across pipeline and data parallel axes.

```python
np.random.seed(42)
total_steps = MICRO_BATCHES + PP_SIZE

w_np = np.random.randn(PP_SIZE, DIM, DIM).astype(np.float32)
b_np = np.random.randn(PP_SIZE, DIM).astype(np.float32)
x_np = np.random.randn(MICRO_BATCHES, MICRO_BATCH_SIZE, DIM).astype(np.float32)
y_np = np.random.randn(MICRO_BATCHES, MICRO_BATCH_SIZE, DIM).astype(np.float32)

# Weights: sharded on "pp", replicated on "dp"
w_spec = [DimSpec.from_raw("pp"), None, None]
b_spec = [DimSpec.from_raw("pp"), None]
w_sharded = ops.shard(nb.Tensor.from_dlpack(w_np), mesh, w_spec).realize()
b_sharded = ops.shard(nb.Tensor.from_dlpack(b_np), mesh, b_spec).realize()

# Data: sharded on "dp" (batch dim), replicated on "pp"
x_padded_np = np.concatenate(
    [x_np, np.zeros((PP_SIZE, MICRO_BATCH_SIZE, DIM), dtype=np.float32)], axis=0
)
x_spec = [None, DimSpec.from_raw("dp"), None]
x_sharded = ops.shard(nb.Tensor.from_dlpack(x_padded_np), mesh, x_spec).realize()
y_sharded = ops.shard(nb.Tensor.from_dlpack(y_np), mesh, x_spec).realize()

# Pipeline state: sharded on both axes
state_spec = [DimSpec.from_raw("pp"), DimSpec.from_raw("dp"), None]
state_sharded = ops.shard(
    nb.zeros((PP_SIZE, MICRO_BATCH_SIZE, DIM)), mesh, state_spec
).realize()

# Stage-0 mask
mask_np = np.eye(PP_SIZE, 1).reshape(PP_SIZE, 1, 1).astype(bool)
mask_spec = [DimSpec.from_raw("pp"), None, None]
mask_sharded = ops.shard(nb.Tensor.from_dlpack(mask_np), mesh, mask_spec).realize()

print(f"Weights: {w_sharded.shape} sharded on 'pp'")
print(f"Inputs:  {x_sharded.shape} sharded on 'dp'")
```

```text
Weights: [Dim(4), Dim(16), Dim(16)] sharded on 'pp'
Inputs:  [Dim(8), Dim(4), Dim(16)] sharded on 'dp'
```

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

### LoRA Training Output

Source: https://www.nablaml.com/examples/11_lora_finetuning.html

Example console output resulting from the LoRA training process.

```text
Initial loss: 32.337105
Step  50: loss = 0.199841
LoRA training time: 1.5058 s
LoRA avg step:      25.097 ms/step
Final loss:   0.068891
✅ LoRA adapter trained successfully
```

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

### Completion Confirmation

Source: https://www.nablaml.com/_sources/examples/04a_mlp_training_pytorch.ipynb

A simple print statement used to signal the successful completion of the training example.

```python
print("\n✅ Example 04a completed!")
```

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

### Compiled Functional Training Setup

Source: https://www.nablaml.com/_sources/examples/06a_transformer_pytorch.ipynb

Defines a functional training step using @nb.compile and value_and_grad, which is required for compiled execution instead of imperative .backward().

```python
import time

# Fresh model + functional optimizer state for compiled run
model_c = TransformerClassifier(
    vocab_size=vocab_size, d_model=d_model, num_heads=num_heads,
    num_layers=num_layers, num_classes=num_classes,
    max_len=seq_len, dim_feedforward=64,
)
# eval() BEFORE adamw_init so both share _training=False pytree structure
model_c.eval()
opt_state_c = nb.nn.optim.adamw_init(model_c)


def loss_fn_for_compile(model, tokens, targets):
    logits = model(tokens)
    return nb.nn.functional.cross_entropy_loss(logits, targets)


@nb.compile
def compiled_step(model, opt_state, tokens, targets):
    loss, grads = nb.value_and_grad(loss_fn_for_compile, argnums=0)(
        model, tokens, targets
    )
    model, opt_state = nb.nn.optim.adamw_update(model, grads, opt_state, lr=1e-3)
    return model, opt_state, loss


def eager_step(model, opt_state, tokens, targets):
    loss, grads = nb.value_and_grad(loss_fn_for_compile, argnums=0)(
        model, tokens, targets
    )
    model, opt_state = nb.nn.optim.adamw_update(model, grads, opt_state, lr=1e-3)
    return model, opt_state, loss
```

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

### Automatic Vectorization with vmap

Source: https://www.nablaml.com/_sources/examples/05_transforms_and_compile.ipynb

Transform a function operating on single examples into one that handles batches automatically.

```python
def single_dot(x, y):
    """Dot product of two vectors (no batch dimension)."""
    return nb.reduce_sum(x * y)

# Without vmap: manual loop
x_batch = nb.uniform((5, 3))
y_batch = nb.uniform((5, 3))

# With vmap: automatic vectorization!
batched_dot = nb.vmap(single_dot, in_axes=(0, 0))
result = batched_dot(x_batch, y_batch)
print(f"Batched dot products (5 pairs of 3D vectors):")
print(result)
print(f"Shape: {result.shape}")
```

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

### arange(start, stop, step, *, dtype, device)

Source: https://www.nablaml.com/_sources/api/ops/creation/index.md

Returns a 1-D tensor with evenly spaced values in the interval [start, stop).

```APIDOC
## arange(start: int = 0, stop: int | None = None, step: int = 1, *, dtype: DType | None = None, device: Device | None = None)

### Description
Returns a 1-D tensor with evenly spaced values in [start, stop). When called with a single positional argument, it is treated as stop and start defaults to 0.

### Parameters
- **start** (int) - Optional - Start of the interval (inclusive). Default: 0.
- **stop** (int | None) - Optional - End of the interval (exclusive). If None, start is used as stop and start becomes 0.
- **step** (int) - Optional - Spacing between values. Default: 1.
- **dtype** (DType | None) - Optional - Element dtype. Uses the current default if None.
- **device** (Device | None) - Optional - Target device. Uses the current default if None.

### Returns
1-D tensor of length ceil((stop - start) / step).
```

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

### Perform QLoRA Fine-Tuning

Source: https://www.nablaml.com/_sources/index.md

Demonstrates quantizing weights to NF4 and initializing a LoRA adapter for fine-tuning.

```python
import nabla as nb

# Quantize frozen weights to NF4
qweight = nb.nn.finetune.quantize_nf4(frozen_weight, block_size=64)

# Initialize LoRA adapter
lora_params = nb.nn.finetune.init_lora_adapter(frozen_weight, rank=8)

# Forward pass with QLoRA
def loss_fn(adapter, batch_x, batch_y):
    pred = nb.nn.finetune.qlora_linear(
        batch_x, qweight, adapter, alpha=16.0
    )
    return nb.mean((pred - batch_y) ** 2)

# Compute gradients
loss, grads = nb.value_and_grad(loss_fn)(lora_params, x, y)
```

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

### Initialize and Train QLoRA Adapter

Source: https://www.nablaml.com/_sources/examples/11_lora_finetuning.ipynb

Initializes a LoRA adapter, defines a MSE loss function for quantized linear layers, and performs iterative training.

```python
# Fresh adapter for QLoRA
qlora_params = nb.nn.finetune.init_lora_adapter(frozen_weight, rank=RANK, init_std=0.01)
qopt_state = nb.nn.optim.adamw_init(qlora_params)

def qlora_loss_fn(adapter, batch_x, batch_y):
    """MSE loss with QLoRA-adapted linear layer."""
    pred = nb.nn.finetune.qlora_linear(
        batch_x, qweight, adapter, alpha=ALPHA, compute_dtype=nb.DType.float32
    )
    diff = pred - batch_y
    return nb.mean(diff * diff)

q_initial_loss = float(qlora_loss_fn(qlora_params, x, y).to_numpy())
print(f"QLoRA initial loss: {q_initial_loss:.6f}")

qlora_start = time.perf_counter()
for step in range(STEPS):
    loss, qlora_params, qopt_state = train_step(qlora_loss_fn, qlora_params, qopt_state, x, y)
    if (step + 1) % 50 == 0:
        print(f"Step {step + 1:>3d}: loss = {float(loss.to_numpy()):.6f}")
qlora_elapsed = time.perf_counter() - qlora_start
print(f"QLoRA training time: {qlora_elapsed:.4f} s")
print(f"QLoRA avg step:      {(qlora_elapsed / max(1, STEPS)) * 1000.0:.3f} ms/step")

q_final_loss = float(qlora_loss_fn(qlora_params, x, y).to_numpy())
print(f"QLoRA final loss:   {q_final_loss:.6f}")
assert q_final_loss < q_initial_loss, "QLoRA training did not reduce loss"
print("✅ QLoRA adapter trained successfully")
```

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

### Get shard shape

Source: https://www.nablaml.com/api/core/tensor/index.html

Returns the shape of a specific shard.

```python
def shard_shape(self, shard_idx: 'int' = 0) -> 'graph.Shape':
```

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

### Get number of elements (alias)

Source: https://www.nablaml.com/api/core/tensor/index.html

PyTorch-style alias for num_elements().

```python
def numel(self) -> 'int':
```

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

### Get number of elements

Source: https://www.nablaml.com/api/core/tensor/index.html

Returns the total number of elements in the tensor.

```python
def num_elements(self) -> 'int':
```

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

### Get sharding rule template

Source: https://www.nablaml.com/_sources/api/ops/custom/index.md

Retrieves the template for the sharding rule.

```python
def get_sharding_rule_template(self) -> 'Any':
```

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

### Shard Data and Initialize Pipeline

Source: https://www.nablaml.com/_sources/examples/07_mlp_pipeline_parallel.ipynb

Configures weight sharding across stages and prepares zero-padded inputs for pipeline warm-up.

```python
np.random.seed(42)

# Random weights (one per stage), inputs, and targets
w_np = np.random.randn(STAGES, DIM, DIM).astype(np.float32)
b_np = np.random.randn(STAGES, DIM).astype(np.float32)
x_np = np.random.randn(MICRO_BATCHES, MICRO_BATCH_SIZE, DIM).astype(np.float32)
y_np = np.random.randn(MICRO_BATCHES, MICRO_BATCH_SIZE, DIM).astype(np.float32)

# Shard weights: first axis → "stage" mesh axis
w_spec = [DimSpec.from_raw(d) for d in P("stage", None, None)]
b_spec = [DimSpec.from_raw(d) for d in P("stage", None)]

w_sharded = ops.shard(nb.Tensor.from_dlpack(w_np), mesh, w_spec)
b_sharded = ops.shard(nb.Tensor.from_dlpack(b_np), mesh, b_spec)

# Pad inputs with STAGES zero-slices for pipeline warm-up
padding = np.zeros((STAGES, MICRO_BATCH_SIZE, DIM), dtype=np.float32)
x_padded = nb.Tensor.from_dlpack(np.concatenate([x_np, padding], axis=0))
y_nb = nb.Tensor.from_dlpack(y_np)

# Initial pipeline state: zeros on each stage
state_sharded = ops.shard(
    nb.zeros((STAGES, MICRO_BATCH_SIZE, DIM)), mesh, w_spec
)

# Stage-0 mask for injecting fresh inputs
mask_np = np.eye(STAGES, 1).reshape(STAGES, 1, 1).astype(bool)
mask_0 = ops.shard(nb.Tensor.from_dlpack(mask_np), mesh, w_spec)

nb.realize_all(w_sharded, b_sharded, state_sharded, mask_0)
print(f"Weights sharded: {w_sharded.shape}, Inputs padded: {x_padded.shape}")
```

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

### Get size or shape

Source: https://www.nablaml.com/api/core/tensor/index.html

Returns the shape or the size of a specific dimension (PyTorch style).

```python
def size(self, dim: 'int | None' = None) -> 'graph.Shape | int':
```

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

### Initialize Nabla Environment

Source: https://www.nablaml.com/_sources/examples/04a_mlp_training_pytorch.ipynb

Import necessary libraries and initialize the Nabla environment.

```python
import time
import numpy as np

import nabla as nb

print("Nabla MLP Training — PyTorch-style")
```

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

### Get physical local shape

Source: https://www.nablaml.com/api/core/tensor/index.html

Returns the physical local shape for a specific shard.

```python
def physical_local_shape(self, shard_idx: 'int' = 0) -> 'graph.Shape | None':
```

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

### Shard Weights and Prepare Inputs

Source: https://www.nablaml.com/_sources/examples/09_mlp_pipeline_inference.ipynb

Configures the device mesh, shards weights across stages, and pads inputs for the pipeline.

```python
mesh = DeviceMesh("pp", (STAGES,), ("stage",))
np.random.seed(42)

w_np = np.random.randn(STAGES, DIM, DIM).astype(np.float32)
x_np = np.random.randn(MICRO_BATCHES, MICRO_BATCH_SIZE, DIM).astype(np.float32)

# Shard weights across stages
w_spec = [DimSpec.from_raw(d) for d in P("stage", None, None)]
w_sharded = ops.shard(nb.Tensor.from_dlpack(w_np), mesh, w_spec).realize()

# Pad inputs for pipeline warm-up
padding = np.zeros((STAGES, MICRO_BATCH_SIZE, DIM), dtype=np.float32)
x_padded = nb.Tensor.from_dlpack(np.concatenate([x_np, padding], axis=0))

# Initial state and stage-0 mask
state_sharded = ops.shard(
    nb.Tensor.from_dlpack(np.zeros((STAGES, MICRO_BATCH_SIZE, DIM), dtype=np.float32)),
    mesh, w_spec
).realize()

mask_np = np.eye(STAGES, 1).reshape(STAGES, 1, 1).astype(bool)
mask_0 = ops.shard(nb.Tensor.from_dlpack(mask_np), mesh, w_spec).realize()

print(f"Mesh: {mesh}")
print(f"Weights: {w_sharded.shape}, Inputs: {x_padded.shape}")
```

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

### Initialize Nabla Tracing

Source: https://www.nablaml.com/_sources/examples/03_graph_tracing.ipynb

Import necessary libraries and the low-level tracing API.

```python
import numpy as np

import nabla as nb
from nabla.core import trace  # low-level tracing API

print("Nabla graph tracing example")
```

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

### Initialize and Train LoRA Adapter

Source: https://www.nablaml.com/examples/11_lora_finetuning.html

Initializes a LoRA adapter, defines the training loop, and executes the optimization process.

```python
lora_params = nb.nn.finetune.init_lora_adapter(frozen_weight, rank=RANK, init_std=0.01)
opt_state = nb.nn.optim.adamw_init(lora_params)

def lora_loss_fn(adapter, batch_x, batch_y):
    """MSE loss with LoRA-adapted linear layer."""
    pred = nb.nn.finetune.lora_linear(batch_x, frozen_weight, adapter, alpha=ALPHA)
    diff = pred - batch_y
    return nb.mean(diff * diff)

def train_step(loss_fn, adapter, optimizer_state, batch_x, batch_y):
    """One training step: forward + backward + AdamW update."""
    loss, grads = nb.value_and_grad(loss_fn, argnums=0, realize=False)(
        adapter, batch_x, batch_y
    )
    new_adapter, new_state = nb.nn.optim.adamw_update(
        adapter, grads, optimizer_state, lr=LR, weight_decay=0.0
    )
    # Batch-realize all tensors
    to_realize = [loss]
    to_realize.extend(t for t in nb.tree_leaves(grads) if isinstance(t, nb.Tensor))
    to_realize.extend(t for t in nb.tree_leaves(new_adapter) if isinstance(t, nb.Tensor))
    to_realize.extend(t for t in nb.tree_leaves(new_state) if isinstance(t, nb.Tensor))
    nb.realize_all(*to_realize)
    return loss, new_adapter, new_state

initial_loss = float(lora_loss_fn(lora_params, x, y).to_numpy())
print(f"Initial loss: {initial_loss:.6f}")

lora_start = time.perf_counter()
for step in range(STEPS):
    loss, lora_params, opt_state = train_step(lora_loss_fn, lora_params, opt_state, x, y)
    if (step + 1) % 50 == 0:
        print(f"Step {step + 1:>3d}: loss = {float(loss.to_numpy()):.6f}")
lora_elapsed = time.perf_counter() - lora_start
print(f"LoRA training time: {lora_elapsed:.4f} s")
print(f"LoRA avg step:      {(lora_elapsed / max(1, STEPS)) * 1000.0:.3f} ms/step")

final_loss = float(lora_loss_fn(lora_params, x, y).to_numpy())
print(f"Final loss:   {final_loss:.6f}")
assert final_loss < initial_loss, "LoRA training did not reduce loss"
print("✅ LoRA adapter trained successfully")
```

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

### Import Dependencies

Source: https://www.nablaml.com/_sources/examples/10_jax_comparison_compiled.ipynb

Initializes the environment and attempts to import JAX for comparison.

```python
import time

import numpy as np

import nabla as nb

# Try to import JAX
try:
    import jax
    import jax.numpy as jnp
    from jax import grad, jit

    HAS_JAX = True
except ImportError:
    HAS_JAX = False
```

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

### Initialize and Train LoRA Adapter

Source: https://www.nablaml.com/_sources/examples/11_lora_finetuning.ipynb

Initializes a LoRA adapter and defines the training step using AdamW optimization.

```python
lora_params = nb.nn.finetune.init_lora_adapter(frozen_weight, rank=RANK, init_std=0.01)
opt_state = nb.nn.optim.adamw_init(lora_params)

def lora_loss_fn(adapter, batch_x, batch_y):
    """MSE loss with LoRA-adapted linear layer."""
    pred = nb.nn.finetune.lora_linear(batch_x, frozen_weight, adapter, alpha=ALPHA)
    diff = pred - batch_y
    return nb.mean(diff * diff)

def train_step(loss_fn, adapter, optimizer_state, batch_x, batch_y):
    """One training step: forward + backward + AdamW update."""
    loss, grads = nb.value_and_grad(loss_fn, argnums=0, realize=False)(
        adapter, batch_x, batch_y
    )
    new_adapter, new_state = nb.nn.optim.adamw_update(
        adapter, grads, optimizer_state, lr=LR, weight_decay=0.0
    )
```

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

### Slice Tensor

Source: https://www.nablaml.com/_sources/api/ops/view/index.md

Extracts a rectangular slice from the input tensor based on start indices and dimensions.

```python
def slice_tensor(x: 'Tensor', start: 'Any', size: 'Any') -> 'Tensor':
```

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

### Get total number of shards

Source: https://www.nablaml.com/_sources/api/core/sharding/index.md

Retrieves the total count of shards defined by a sharding specification.

```python
def get_num_shards(sharding: nabla.core.sharding.spec.ShardingSpec) -> int:
```

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

### Compiled vs Eager Training Steps

Source: https://www.nablaml.com/_sources/examples/06b_transformer_jax.ipynb

Demonstrates how to use the @nb.compile decorator to JIT-compile a training step compared to an eager baseline.

```python
@nb.compile
def compiled_step(params, opt_state, tokens, targets):
    loss, grads = nb.value_and_grad(loss_fn, argnums=0)(params, tokens, targets)
    params, opt_state = nb.nn.optim.adamw_update(
        params, grads, opt_state, lr=1e-3
    )
    return params, opt_state, loss


def eager_step(params, opt_state, tokens, targets):
    """Same logic without @nb.compile — serves as eager baseline."""
    loss, grads = nb.value_and_grad(loss_fn, argnums=0)(params, tokens, targets)
    params, opt_state = nb.nn.optim.adamw_update(
        params, grads, opt_state, lr=1e-3
    )
    return params, opt_state, loss
```

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

### Executing the Compiled Training Loop

Source: https://www.nablaml.com/_sources/examples/05_transforms_and_compile.ipynb

Initializes the model and runs a training loop, performing a warmup step before measuring performance.

```python
# Setup data and model
np.random.seed(0)
X = nb.Tensor.from_dlpack(np.random.randn(100, 4).astype(np.float32))
y = nb.Tensor.from_dlpack(np.random.randn(100, 1).astype(np.float32))

num_steps = 60

# --- Compiled ---
model = TinyMLP()
opt_state = nb.nn.optim.adamw_init(model)

# First call is warmup (trace + compile)
model, opt_state, _ = train_step(model, opt_state, X, y)

print(f"Compiled training loop:")
print(f"{'Step':<8} {'Loss':<12}")
print("-" * 22)

train_start = time.perf_counter()
for step in range(num_steps):
    model, opt_state, loss = train_step(model, opt_state, X, y)
    if (step + 1) % 10 == 0:
        print(f"{step + 1:<8} {loss.item():<12.6f}")
train_elapsed = time.perf_counter() - train_start
```

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

### Importing Required Modules

Source: https://www.nablaml.com/_sources/examples/12_custom_mojo_kernels.ipynb

Initial imports for setting up custom kernel integration in Nabla.

```python
from pathlib import Path
import numpy as np

import nabla as nb
from max.graph import TensorValue
from nabla.ops import UnaryOperation, call_custom_kernel
```

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

### Initialize Dataset

Source: https://www.nablaml.com/_sources/examples/10_jax_comparison_compiled.ipynb

Generates synthetic data and converts it into Nabla tensors.

```python
np.random.seed(42)
n_samples = 500
n_steps = 60

X_np = np.linspace(0, 1, n_samples).reshape(-1, 1).astype(np.float32)
y_np = (np.sin(8 * np.pi * X_np) + 1) / 2.0

X = nb.Tensor.from_dlpack(X_np)
y = nb.Tensor.from_dlpack(y_np)

print(f"Dataset: {n_samples} samples, fitting (sin(8π·x) + 1)/2")
```

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

### Create a 1-D tensor with arange

Source: https://www.nablaml.com/_sources/api/ops/creation/index.md

Generates a 1-D tensor with evenly spaced values in the interval [start, stop).

```python
def arange(start: 'int' = 0, stop: 'int | None' = None, step: 'int' = 1, *, dtype: 'DType | None' = None, device: 'Device | None' = None):
```

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

### Initialize Parameters as Pytree

Source: https://www.nablaml.com/_sources/examples/04b_mlp_training_jax.ipynb

Create a nested dictionary structure to hold model weights and biases instead of using a class.

```python
def init_mlp_params(in_dim: int, hidden_dim: int, out_dim: int) -> dict:
    """Initialize MLP parameters as a pytree (nested dict)."""
    params = {
        "fc1": {
            "weight": xavier_normal((in_dim, hidden_dim)),
            "bias": nb.zeros((1, hidden_dim)),
        },
        "fc2": {
            "weight": xavier_normal((hidden_dim, out_dim)),
            "bias": nb.zeros((1, out_dim)),
        },
    }
    return params


params = init_mlp_params(4, 32, 1)
print("Parameter shapes:")
for name, layer in params.items():
    for pname, p in layer.items():
        print(f"  {name}.{pname}: {p.shape}")
```

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

### Get physical local shape as ints

Source: https://www.nablaml.com/api/core/tensor/index.html

Returns the shape as a tuple of integers for a specific shard, avoiding object creation.

```python
def physical_local_shape_ints(self, shard_idx: 'int' = 0) -> 'tuple[int, ...] | None':
```

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

### Integrate Mojo Custom Kernels

Source: https://www.nablaml.com/_sources/get_started.md

Shows how to define a high-performance Mojo kernel and invoke it from Python using Nabla's custom operation interface.

```mojo
@compiler.register("my_kernel")
struct MyKernel:
    @staticmethod
    def execute[target: StaticString](
        output: OutputTensor,
        x: InputTensor[dtype = output.dtype, rank = output.rank],
        ctx: DeviceContextPtr,
    ):
        @parameter
        fn add_one[W: Int](idx: IndexList[x.rank]) -> SIMD[x.dtype, W]:
            return x.load[W](idx) + 1

        foreach[add_one, target=target](output, ctx)
```

```python
class AddOneOp(nabla.UnaryOperation):
    name = "my_kernel"

    def kernel(self, x, **kwargs):
        # Concise invocation: (func_name, path, inputs, out_types)
        return nabla.call_custom_kernel("my_kernel", "./kernels", x, x.type)

x = nabla.Tensor.constant([1., 2., 3.])
y = AddOneOp()(x)
```

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

### Initialize JAX-style parameters and loss function

Source: https://www.nablaml.com/_sources/examples/05_transforms_and_compile.ipynb

Defines the parameter dictionary structure and the loss function required for JAX-style training.

```python
from nabla.nn.functional import xavier_normal


def init_params():
    return {
        "w1": xavier_normal((4, 16)),
        "b1": nb.zeros((1, 16)),
        "w2": xavier_normal((16, 1)),
        "b2": nb.zeros((1, 1)),
    }


def forward(params, x):
    h = nb.relu(x @ params["w1"] + params["b1"])
    return h @ params["w2"] + params["b2"]


def jax_loss_fn(params, x, y):
    pred = forward(params, x)
    diff = pred - y
    return nb.mean(diff * diff)
```

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

### Manual SGD Implementation

Source: https://www.nablaml.com/_sources/examples/04b_mlp_training_jax.ipynb

Demonstrates manual gradient descent using nb.tree_map to apply updates directly to parameters without an optimizer.

```python
params_sgd = init_mlp_params(4, 32, 1)
sgd_lr = 0.05

print(f"\nManual SGD training:")
print(f"{'Step':<8} {'Loss':<12}")
print("-" * 22)

sgd_steps = 60
sgd_start = time.perf_counter()
for step in range(sgd_steps):
    loss, grads = nb.value_and_grad(loss_fn, argnums=0)(params_sgd, X, y)

    # Manual SGD: params = params - lr * grads
    params_sgd = nb.tree_map(
        lambda p, g: p - sgd_lr * g, params_sgd, grads
    )

    if (step + 1) % 20 == 0:
        print(f"{step + 1:<8} {loss.item():<12.6f}")
sgd_elapsed = time.perf_counter() - sgd_start
print(f"\nManual SGD time: {sgd_elapsed:.4f} s")
print(f"Manual SGD avg step: {(sgd_elapsed / max(1, sgd_steps)) * 1000.0:.3f} ms/step")
```

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

### Configure Pipeline and Device Mesh

Source: https://www.nablaml.com/_sources/examples/07_mlp_pipeline_parallel.ipynb

Defines the pipeline stages and micro-batch configuration, and initializes a 1D device mesh.

```python
# Pipeline configuration
STAGES = 4           # Number of pipeline stages (= devices)
MICRO_BATCHES = 8    # Number of micro-batches to stream through
MICRO_BATCH_SIZE = 4 # Samples per micro-batch
DIM = 16             # Hidden dimension

# Create a 1D device mesh for pipeline parallelism
mesh = DeviceMesh("pp", (STAGES,), ("stage",))
print(f"Device mesh: {STAGES} stages")
print(f"Pipeline: {MICRO_BATCHES} micro-batches × {MICRO_BATCH_SIZE} samples")
```

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

### Execute Custom Mojo Kernel

Source: https://www.nablaml.com/_sources/examples/12_custom_mojo_kernels.ipynb

Demonstrates invoking a custom kernel via a Nabla tensor and verifying the output.

```python
x = nb.Tensor.from_dlpack(np.array([1.0, 2.0, 3.0], dtype=np.float32))
y = add_one(x)
print(f"Input:  {x.to_numpy()}")
print(f"Output: {y.to_numpy()}")  # [2.0, 3.0, 4.0]
```

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

### Implement Pipeline Step

Source: https://www.nablaml.com/_sources/examples/07_mlp_pipeline_parallel.ipynb

Core GPipe logic performing computation, shifting activations via ppermute, and injecting new micro-batches.

```python
def pipeline_step(
    current_state, fresh_input, weight_stack, bias_stack, mask_0, step_fn, perm
):
    """Single GPipe step: compute -> shift -> extract result -> inject input."""
    computed = step_fn(current_state, weight_stack, bias_stack)
    shifted = communication.ppermute(computed, perm)
    # Extract the final stage's output (mask selects stage 0 after the shift)
    res_part = ops.where(mask_0, shifted, ops.zeros_like(shifted))
    result = ops.reduce_sum(res_part, axis=0)
    # Inject the fresh micro-batch at stage 0, pass shifted activations elsewhere
    next_state = ops.where(mask_0, fresh_input, shifted)
    return next_state, result
```

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

### load_finetune_checkpoint

Source: https://www.nablaml.com/_sources/api/nn/finetune/checkpoint/index.md

Loads LoRA parameters and optional optimizer state from a specified checkpoint file.

```APIDOC
## load_finetune_checkpoint

### Description
Load LoRA params and optional optimizer state from checkpoint.

### Signature
`load_finetune_checkpoint(path: 'PathLike', *, lora_template: 'Any', optimizer_template: 'dict[str, Any] | None' = None) -> 'tuple[Any, dict[str, Any] | None, dict[str, Any]]'`

### Parameters
- **path** (PathLike) - Required - The file path to the checkpoint.
- **lora_template** (Any) - Required - The template for loading LoRA parameters.
- **optimizer_template** (dict[str, Any] | None) - Optional - The template for loading the optimizer state.

### Returns
- **tuple** - Returns a tuple containing the loaded LoRA parameters, the optimizer state (if applicable), and the metadata dictionary.
```

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

### Display Training Summary

Source: https://www.nablaml.com/_sources/examples/10_jax_comparison_compiled.ipynb

Prints a summary of training results, including speedup factors, cache hit rates, and JAX JIT comparison status.

```python
print("=" * 70)
print("SUMMARY")
print("=" * 70)
print("✓ MLP training works with compile!")
print("✓ Full pytree parameters (weights + biases) work correctly")
print(
    f"✓ Loss decreases properly: {losses_compiled[0]:.6f} -> {losses_compiled[-1]:.6f}"
)
print(f"✓ {speedup:.2f}x speedup from compilation")
print(f"✓ Cache hit rate: {train_step_compiled.stats.hit_rate:.1f}%")
if HAS_JAX:
    print("✓ Compared against JAX JIT successfully")
```

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

### Print Performance Summary

Source: https://www.nablaml.com/examples/10_jax_comparison_compiled.html

Displays execution time and steps per second for different Nabla and JAX modes, including a comparison of speedups and verification of loss consistency.

```python
print("=" * 60)
print("PERFORMANCE SUMMARY")
print("=" * 60)
print(f"Nabla @nb.compile:  {elapsed_compiled:.4f}s  ({n_steps / elapsed_compiled:.1f} steps/sec)")
print(f"Nabla Eager:        {elapsed_eager:.4f}s  ({n_steps / elapsed_eager:.1f} steps/sec)")
print(f"Nabla Eager (MAX):  {elapsed_eager_max:.4f}s  ({n_steps / elapsed_eager_max:.1f} steps/sec)")

if HAS_JAX:
    print(f"JAX @jit:           {elapsed_jax:.4f}s  ({n_steps / elapsed_jax:.1f} steps/sec)")
    speedup_vs_jax = elapsed_jax / elapsed_compiled
    if speedup_vs_jax > 1:
        print(f"\n🚀 Nabla compiled is {speedup_vs_jax:.2f}x faster than JAX JIT")
    else:
        print(f"\nJAX JIT is {1 / speedup_vs_jax:.2f}x faster than Nabla compiled")

speedup = elapsed_eager / elapsed_compiled
print(f"\nCompile speedup over eager: {speedup:.2f}x")

# Verify correctness across modes
loss_diff = abs(losses_compiled[-1] - losses_eager[-1])
print(f"Loss match (compiled vs eager): {'✅' if loss_diff < 1e-4 else '⚠️'} diff={loss_diff:.8f}")
```

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

### Initialize Model and Optimizer

Source: https://www.nablaml.com/_sources/examples/06a_transformer_pytorch.ipynb

Instantiates the Transformer classifier and configures the AdamW optimizer while ensuring the model is in training mode.

```python
model = TransformerClassifier(
    vocab_size=vocab_size,
    d_model=d_model,
    num_heads=num_heads,
    num_layers=num_layers,
    num_classes=num_classes,
    max_len=seq_len,
    dim_feedforward=64,
)

n_params = sum(p.numel() for p in model.parameters())
print(f"Model: {num_layers} encoder layers, d_model={d_model}, heads={num_heads}")
print(f"Total trainable parameters: {n_params}")

# Create optimizer while model is in train mode (default _training=True)
model.train()
optimizer = nb.nn.optim.AdamW(model, lr=1e-3)
print(f"Optimizer: AdamW (lr={optimizer.lr})")
```

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

### Compiled JAX-style training step

Source: https://www.nablaml.com/_sources/examples/05_transforms_and_compile.ipynb

Implements the training loop using @nb.compile to fuse the gradient calculation and optimizer update, compared against an eager baseline.

```python
@nb.compile
def jax_train_step(params, opt_state, x, y):
    loss, grads = nb.value_and_grad(jax_loss_fn, argnums=0)(params, x, y)
    params, opt_state = nb.nn.optim.adamw_update(
        params, grads, opt_state, lr=1e-2
    )
    return params, opt_state, loss


def eager_jax_train_step(params, opt_state, x, y):
    """Same logic without @nb.compile — eager baseline."""
    loss, grads = nb.value_and_grad(jax_loss_fn, argnums=0)(params, x, y)
    params, opt_state = nb.nn.optim.adamw_update(
        params, grads, opt_state, lr=1e-2
    )
    return params, opt_state, loss


# --- Compiled ---
params = init_params()
opt_state = nb.nn.optim.adamw_init(params)
params, opt_state, _ = jax_train_step(params, opt_state, X, y)  # warmup

print(f"Compiled JAX-style training:")
print(f"{'Step':<8} {'Loss':<12}")
print("-" * 22)

jax_steps = 60
jax_start = time.perf_counter()
for step in range(jax_steps):
    params, opt_state, loss = jax_train_step(params, opt_state, X, y)
    if (step + 1) % 10 == 0:
        print(f"{step + 1:<8} {loss.item():<12.6f}")
jax_elapsed = time.perf_counter() - jax_start
compiled_jax_ms = (jax_elapsed / jax_steps) * 1000.0

print(f"\nCompiled cache stats: {jax_train_step.stats}")
print(f"Compiled JAX avg step: {compiled_jax_ms:.3f} ms/step")

# --- Eager baseline ---
params_e = init_params()
opt_state_e = nb.nn.optim.adamw_init(params_e)
params_e, opt_state_e, _ = eager_jax_train_step(params_e, opt_state_e, X, y)
```

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

### Define a PartitionSpec

Source: https://www.nablaml.com/_sources/api/core/sharding/index.md

Creates a JAX-compatible partition specification.

```python
class PartitionSpec(*args):
```

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

### Initialize LoRA Adapter

Source: https://www.nablaml.com/api/nn/finetune/lora/index.html

Initializes LoRA adapter matrices A and B for a 2D weight tensor.

```python
def init_lora_adapter(weight: 'Tensor', rank: 'int', init_std: 'float' = 0.01, dtype: 'DType | None' = None) -> 'dict[str, Tensor]':
```

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

### Define a P

Source: https://www.nablaml.com/_sources/api/core/sharding/index.md

Creates a JAX-compatible partition specification alias.

```python
class P(*args):
```

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

### Compare Compiled and Eager Transformer Training

Source: https://www.nablaml.com/_sources/examples/06b_transformer_jax.ipynb

Initializes model parameters and compares the execution speed of a JIT-compiled training step against an eager functional baseline.

```python
# Fresh parameters for compiled training
params2 = init_transformer(
    vocab_size=vocab_size, d_model=d_model, num_layers=num_layers,
    num_heads=num_heads, num_classes=num_classes, max_len=seq_len, dim_ff=dim_ff,
)
opt_state2 = nb.nn.optim.adamw_init(params2)

# First call is warmup (trace + compile)
params2, opt_state2, _ = compiled_step(params2, opt_state2, token_ids, labels)

print(f"Compiled training:")
print(f"{'Step':<8} {'Loss':<12}")
print("-" * 22)

compiled_steps = num_epochs
compiled_start = time.perf_counter()
for step in range(compiled_steps):
    params2, opt_state2, loss = compiled_step(params2, opt_state2, token_ids, labels)
    if (step + 1) % 10 == 0:
        print(f"{step + 1:<8} {loss.item():<12.4f}")
compiled_elapsed = time.perf_counter() - compiled_start
compiled_step_ms = (compiled_elapsed / max(1, compiled_steps)) * 1000.0

print(f"\nCompiled cache stats: {compiled_step.stats}")
print(f"Compiled avg step:   {compiled_step_ms:.2f} ms")

# --- Eager functional baseline (same math, no @nb.compile) ---
params_e = init_transformer(
    vocab_size=vocab_size, d_model=d_model, num_layers=num_layers,
    num_heads=num_heads, num_classes=num_classes, max_len=seq_len, dim_ff=dim_ff,
)
opt_state_e = nb.nn.optim.adamw_init(params_e)
params_e, opt_state_e, _ = eager_step(params_e, opt_state_e, token_ids, labels)  # warmup

eager_start = time.perf_counter()
for step in range(compiled_steps):
    params_e, opt_state_e, loss_e = eager_step(params_e, opt_state_e, token_ids, labels)
eager_elapsed = time.perf_counter() - eager_start
eager_step_ms = (eager_elapsed / max(1, compiled_steps)) * 1000.0

speedup = eager_step_ms / max(compiled_step_ms, 1e-9)
print(f"Eager avg step:      {eager_step_ms:.2f} ms")
print(f"Compiled speedup:    {speedup:.1f}x")
```

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

### Define PartitionSpec and P

Source: https://www.nablaml.com/api/core/sharding/index.html

JAX-compatible specifications for partitioning tensors across devices.

```python
class PartitionSpec(*args):
```

```python
class P(*args):
```

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

### Module methods definitions

Source: https://www.nablaml.com/api/nn/modules/index.html

Signatures for standard methods available on Module instances.

```python
def backward(self, loss: 'Tensor', gradient: 'Tensor | None' = None, retain_graph: 'bool' = False, create_graph: 'bool' = False, *, realize_grads: 'bool | None' = None) -> 'None':

```

```python
def buffers(self) -> 'Iterator[Tensor]':

```

```python
def eval(self) -> 'Module':

```

```python
def extra_repr(self) -> 'str':

```

```python
def forward(self, *args: 'Any', **kwargs: 'Any') -> 'Any':

```

```python
def load_state_dict(self, state_dict: 'OrderedDict[str, Tensor]') -> 'None':

```

```python
def modules(self) -> 'Iterator[Module]':

```

```python
def named_buffers(self, prefix: 'str' = '') -> 'Iterator[tuple[str, Tensor]]':

```

```python
def named_parameters(self, prefix: 'str' = '') -> 'Iterator[tuple[str, Tensor]]':

```

```python
def parameters(self) -> 'Iterator[Tensor]':

```

```python
def register_buffer(self, name: 'str', tensor: 'Tensor | None') -> 'None':

```

```python
def state_dict(self) -> 'OrderedDict[str, Tensor]':

```

```python
def train(self) -> 'Module':

```

```python
def zero_grad(self) -> 'None':

```
Tokens:
Raw