1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95

import torch
import torch.nn as nn
import torch.nn.functional as F
 
# ============ Understanding PyTorch's Graph Building ============
 
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = torch.tensor([4.0, 5.0, 6.0], requires_grad=True)
 
# Each operation creates a new tensor with grad_fn
z = x * y  # z has grad_fn=<MulBackward0>
print(f"z.grad_fn: {z.grad_fn}")
 
w = z.sum()  # w has grad_fn=<SumBackward0>
print(f"w.grad_fn: {w.grad_fn}")
 
# grad_fn contains references to input grad_fns
print(f"w's next functions: {w.grad_fn.next_functions}")
# Shows the MulBackward0 that created z
 
# ============ Custom Autograd Function ============
 
class MySigmoid(torch.autograd.Function):
    """
    Custom autograd function demonstrating how PyTorch
    implements forward/backward for primitives.
    """
    
    @staticmethod
    def forward(ctx, x):
        """
        Forward pass.
        ctx is a context object for passing info to backward.
        """
        result = 1 / (1 + torch.exp(-x))
        # Save for backward - this is what gets cached
        ctx.save_for_backward(result)
        return result
    
    @staticmethod
    def backward(ctx, grad_output):
        """
        Backward pass.
        grad_output is the upstream gradient (dL/d(sigmoid output))
        Returns gradient w.r.t. each input.
        """
        result, = ctx.saved_tensors
        # sigmoid'(x) = sigmoid(x) * (1 - sigmoid(x))
        grad_input = grad_output * result * (1 - result)
        return grad_input
 
# Using custom function
my_sigmoid = MySigmoid.apply
x = torch.randn(5, requires_grad=True)
y = my_sigmoid(x)
loss = y.sum()
loss.backward()
print(f"Custom sigmoid gradient: {x.grad}")
 
# ============ Gradient Accumulation and zeroing ============
 
# Gradients ACCUMULATE by default - crucial for RNNs, gradient accumulation
x = torch.tensor([1.0], requires_grad=True)
 
# First backward
(x ** 2).backward()
print(f"After first backward: {x.grad}")  # 2.0
 
# Second backward (without zeroing)
(x ** 3).backward()
print(f"After second backward: {x.grad}")  # 2.0 + 3.0 = 5.0 (accumulated!)
 
# Must zero gradients before each new computation
x.grad.zero_()
(x ** 2).backward()
print(f"After zeroing and new backward: {x.grad}")  # 2.0
 
# ============ Gradient Checkpointing ============
from torch.utils.checkpoint import checkpoint
 
class LargeModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([
            nn.Linear(1024, 1024) for _ in range(20)
        ])
    
    def forward(self, x):
        for i, layer in enumerate(self.layers):
            # Checkpoint every 4th layer to save memory
            if i % 4 == 0:
                x = checkpoint(lambda l, x: F.relu(l(x)), layer, x)
            else:
                x = F.relu(layer(x))
        return x

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96

import tensorflow as tf
from tensorflow import keras
 
# ============ Eager vs Graph Mode ============
 
# Eager execution (default in TF2)
x = tf.constant([1.0, 2.0, 3.0])
y = x ** 2  # Executes immediately
print(f"Eager result: {y}")  # tf.Tensor([1. 4. 9.], ...)
 
# Graph mode with tf.function
@tf.function
def square_fn(x):
    return x ** 2
 
# First call: traces the function and compiles graph
result = square_fn(tf.constant([1.0, 2.0]))
# Subsequent calls: uses compiled graph (fast)
result = square_fn(tf.constant([3.0, 4.0]))
 
# ============ GradientTape for Autodiff ============
 
x = tf.Variable([3.0, 4.0])
 
# Record operations within tape context
with tf.GradientTape() as tape:
    y = x ** 2
    z = tf.reduce_sum(y)
 
# Compute gradient (this is the backward pass)
grad = tape.gradient(z, x)
print(f"Gradient: {grad}")  # [6. 8.] = 2*x
 
# Persistent tape for multiple gradient computations
x = tf.Variable([3.0])
with tf.GradientTape(persistent=True) as tape:
    y = x ** 2
    z = x ** 3
 
grad_y = tape.gradient(y, x)  # 6.0
grad_z = tape.gradient(z, x)  # 27.0
del tape  # Must delete persistent tape
 
# ============ Higher-Order Gradients ============
 
x = tf.Variable(3.0)
 
with tf.GradientTape() as outer_tape:
    with tf.GradientTape() as inner_tape:
        y = x ** 3  # y = x^3
    first_order = inner_tape.gradient(y, x)  # dy/dx = 3x^2 = 27
second_order = outer_tape.gradient(first_order, x)  # d²y/dx² = 6x = 18
 
print(f"First derivative: {first_order}")   # 27.0
print(f"Second derivative: {second_order}")  # 18.0
 
# ============ Understanding tf.function Tracing ============
 
@tf.function
def trace_demo(x):
    print(f"Tracing with input shape: {x.shape}")  # Printed during tracing
    tf.print(f"Executing with value: ", x)  # Printed during execution
    return x + 1
 
# First call with shape (2,) - traces
trace_demo(tf.constant([1.0, 2.0]))
 
# Second call with same shape - reuses trace
trace_demo(tf.constant([3.0, 4.0]))
 
# Third call with different shape - traces AGAIN
trace_demo(tf.constant([1.0, 2.0, 3.0]))
 
# ============ Custom Training Loop ============
 
model = keras.Sequential([
    keras.layers.Dense(64, activation='relu'),
    keras.layers.Dense(10)
])
 
optimizer = keras.optimizers.Adam(learning_rate=1e-3)
loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
 
@tf.function  # Compile for performance
def train_step(x, y):
    with tf.GradientTape() as tape:
        logits = model(x, training=True)
        loss = loss_fn(y, logits)
    
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss
 
# Training would look like:
# for batch_x, batch_y in dataset:
#     loss = train_step(batch_x, batch_y)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109

import jax
import jax.numpy as jnp
from jax import grad, jit, vmap
 
# ============ Basic JAX Operations ============
 
# JAX arrays work like NumPy
x = jnp.array([1.0, 2.0, 3.0])
y = jnp.sin(x) ** 2
print(f"Result: {y}")
 
# ============ Automatic Differentiation ============
 
def f(x):
    """Simple scalar function."""
    return jnp.sin(x) ** 2
 
# grad transforms f into a function that computes df/dx
df = grad(f)
print(f"f(π/4) = {f(jnp.pi/4):.4f}")      # 0.5
print(f"f'(π/4) = {df(jnp.pi/4):.4f}")    # 2*sin(π/4)*cos(π/4) = 1.0
 
# Higher-order gradients via composition
d2f = grad(grad(f))  # Second derivative
print(f"f''(π/4) = {d2f(jnp.pi/4):.4f}")  # cos(2x)*2 at π/4 = 0
 
# ============ Gradients for Multi-Input Functions ============
 
def loss(params, x, y):
    """MSE loss for linear regression."""
    w, b = params
    pred = w * x + b
    return jnp.mean((pred - y) ** 2)
 
# grad with argnums specifies which argument to diff w.r.t.
grad_loss = grad(loss, argnums=0)  # Gradient w.r.t. params
 
params = (2.0, 1.0)  # w=2, b=1
x = jnp.array([1.0, 2.0, 3.0])
y = jnp.array([3.0, 5.0, 7.0])
 
grads = grad_loss(params, x, y)
print(f"Gradients: dw={grads[0]:.4f}, db={grads[1]:.4f}")
 
# ============ JIT Compilation ============
 
def slow_function(x):
    """Function that would be slow without JIT."""
    for _ in range(100):
        x = jnp.sin(x) + jnp.cos(x)
    return x
 
# Without JIT
import time
x = jnp.ones(1000)
 
start = time.time()
result = slow_function(x)
print(f"Without JIT: {time.time() - start:.4f}s")
 
# With JIT - first call compiles
fast_function = jit(slow_function)
_ = fast_function(x)  # Compilation happens here
 
# Subsequent calls are fast
start = time.time()
result = fast_function(x)
print(f"With JIT (after compile): {time.time() - start:.4f}s")
 
# ============ Vectorization with vmap ============
 
def single_example_loss(w, x, y):
    """Loss for a single example."""
    pred = jnp.dot(w, x)
    return (pred - y) ** 2
 
# Manually batching is verbose and error-prone
# vmap transforms to handle batches automatically
 
batch_loss = vmap(single_example_loss, in_axes=(None, 0, 0))
# None means w is not batched; 0 means x and y have batch dim at axis 0
 
w = jnp.array([1.0, 2.0, 3.0])
x_batch = jnp.array([[1.0, 0.0, 0.0],
                      [0.0, 1.0, 0.0],
                      [0.0, 0.0, 1.0]])
y_batch = jnp.array([1.0, 2.0, 3.0])
 
losses = batch_loss(w, x_batch, y_batch)
print(f"Batched losses: {losses}")  # [0. 0. 0.] (perfect predictions)
 
# ============ Composing Transformations ============
 
def model(params, x):
    """Simple MLP."""
    w1, b1, w2, b2 = params
    h = jax.nn.relu(x @ w1 + b1)
    return h @ w2 + b2
 
def loss_fn(params, x, y):
    pred = model(params, x)
    return jnp.mean((pred - y) ** 2)
 
# Compose: JIT-compile the gradient of the loss
fast_grad = jit(grad(loss_fn))
 
# Now fast_grad is a JIT-compiled gradient function
# params = ...
# grads = fast_grad(params, x_batch, y_batch)  # Fast!

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83

# ============ PyTorch 2.0 Compilation ============
import torch
 
model = ...  # Your model
 
# torch.compile analyzes model and generates optimized code
optimized_model = torch.compile(model)
 
# Different modes trade compile time for performance
optimized_model = torch.compile(model, mode="reduce-overhead")  # Lower latency
optimized_model = torch.compile(model, mode="max-autotune")     # Highest throughput
 
# ============ Mixed Precision Training ============
 
# PyTorch AMP (Automatic Mixed Precision)
from torch.cuda.amp import autocast, GradScaler
 
scaler = GradScaler()
 
for batch in dataloader:
    optimizer.zero_grad()
    
    with autocast():  # Automatic float16 where safe
        outputs = model(batch)
        loss = criterion(outputs, targets)
    
    # Scale loss to prevent gradient underflow in float16
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
 
# TensorFlow Mixed Precision
import tensorflow as tf
from tensorflow.keras import mixed_precision
 
# Set global policy
mixed_precision.set_global_policy('mixed_float16')
 
# Model will use float16 for compute, float32 for variables
model = keras.Sequential([...])
model.compile(optimizer='adam', loss='mse')
 
# ============ Memory Optimization ============
 
# Gradient Checkpointing (PyTorch)
from torch.utils.checkpoint import checkpoint_sequential
 
class DeepModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = torch.nn.ModuleList([
            torch.nn.Linear(1024, 1024) for _ in range(50)
        ])
    
    def forward(self, x):
        # Checkpoint every segment of 5 layers
        segments = [self.layers[i:i+5] for i in range(0, 50, 5)]
        return checkpoint_sequential(segments, 5, x)
 
# ============ Data Loading Optimization ============
 
# PyTorch DataLoader optimization
from torch.utils.data import DataLoader
 
loader = DataLoader(
    dataset,
    batch_size=64,
    num_workers=8,           # Parallel data loading
    pin_memory=True,         # Faster CPU→GPU transfer
    prefetch_factor=2,       # Prefetch batches
    persistent_workers=True  # Keep workers alive
)
 
# TensorFlow tf.data optimization
import tensorflow as tf
 
dataset = tf.data.Dataset.from_tensor_slices((x, y))
dataset = (dataset
    .shuffle(buffer_size=10000)
    .batch(64)
    .prefetch(tf.data.AUTOTUNE)  # Overlap data prep with compute
    .cache()                      # Cache in memory after first epoch
)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

# ============ PyTorch DDP (Distributed Data Parallel) ============
 
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
 
def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)
 
def train(rank, world_size):
    setup(rank, world_size)
    
    model = MyModel().to(rank)
    model = DDP(model, device_ids=[rank])
    
    optimizer = torch.optim.Adam(model.parameters())
    
    for batch in dataloader:
        optimizer.zero_grad()
        loss = model(batch).sum()
        loss.backward()  # Gradients synchronized automatically
        optimizer.step()
    
    dist.destroy_process_group()
 
# Launch with torchrun:
# torchrun --nproc_per_node=4 train.py
 
# ============ PyTorch FSDP (Fully Sharded Data Parallel) ============
 
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 
# FSDP shards model parameters across devices
# Enables training models larger than single GPU memory
model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    mixed_precision=MixedPrecision(param_dtype=torch.float16),
)
 
# ============ TensorFlow Distribution Strategy ============
 
import tensorflow as tf
 
# Multi-GPU on single machine
strategy = tf.distribute.MirroredStrategy()
 
with strategy.scope():
    model = create_model()
    model.compile(optimizer='adam', loss='mse')
 
model.fit(dataset, epochs=10)
 
# Multi-machine distributed
strategy = tf.distribute.MultiWorkerMirroredStrategy()
 
# TPU strategy
resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.TPUStrategy(resolver)
 
# ============ JAX Distributed Training ============
 
import jax
from jax import pjit
from jax.sharding import PartitionSpec as P
 
# JAX uses explicit sharding specifications
@partial(pjit, in_axis_resources=(P('data'), P()),
                out_axis_resources=P('data'))
def train_step(batch, params):
    def loss_fn(params):
        return compute_loss(params, batch)
    
    grads = jax.grad(loss_fn)(params)
    # Gradients across devices are automatically synchronized
    return update_params(params, grads)
 
# Model parallelism in JAX
@partial(pjit, in_axis_resources=(None, P('model')),
                out_axis_resources=P('model'))
def forward(x, params):
    return model(params, x)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140

import torch
import torch.nn as nn
from torch.utils.tensorboard import SummaryWriter
 
class TrainingPipeline:
    """
    Robust training pattern incorporating best practices.
    Framework: PyTorch (similar patterns apply to TF/JAX)
    """
    
    def __init__(self, model, train_loader, val_loader, config):
        self.model = model
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.config = config
        
        # Reproducibility
        self._set_seeds(config.seed)
        
        # Device setup
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)
        
        # Optimizer and scheduler
        self.optimizer = torch.optim.AdamW(
            model.parameters(),
            lr=config.learning_rate,
            weight_decay=config.weight_decay
        )
        self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            self.optimizer, T_max=config.epochs
        )
        
        # Mixed precision
        self.scaler = torch.cuda.amp.GradScaler()
        
        # Logging
        self.writer = SummaryWriter(config.log_dir)
        
        # Checkpointing
        self.best_val_loss = float('inf')
        
    def _set_seeds(self, seed):
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        import numpy as np
        np.random.seed(seed)
        import random
        random.seed(seed)
        
    def train_epoch(self, epoch):
        self.model.train()
        total_loss = 0
        
        for batch_idx, (data, target) in enumerate(self.train_loader):
            data, target = data.to(self.device), target.to(self.device)
            
            self.optimizer.zero_grad(set_to_none=True)  # Slightly faster
            
            # Mixed precision forward pass
            with torch.cuda.amp.autocast():
                output = self.model(data)
                loss = nn.functional.cross_entropy(output, target)
            
            # Scaled backward pass
            self.scaler.scale(loss).backward()
            
            # Gradient clipping
            self.scaler.unscale_(self.optimizer)
            torch.nn.utils.clip_grad_norm_(
                self.model.parameters(), 
                self.config.max_grad_norm
            )
            
            self.scaler.step(self.optimizer)
            self.scaler.update()
            
            total_loss += loss.item()
            
            # Log every N steps
            if batch_idx % 100 == 0:
                step = epoch * len(self.train_loader) + batch_idx
                self.writer.add_scalar('train/loss', loss.item(), step)
                
        return total_loss / len(self.train_loader)
    
    def validate(self, epoch):
        self.model.eval()
        val_loss = 0
        correct = 0
        
        with torch.no_grad():
            for data, target in self.val_loader:
                data, target = data.to(self.device), target.to(self.device)
                
                with torch.cuda.amp.autocast():
                    output = self.model(data)
                    val_loss += nn.functional.cross_entropy(output, target).item()
                    pred = output.argmax(dim=1)
                    correct += pred.eq(target).sum().item()
        
        val_loss /= len(self.val_loader)
        accuracy = correct / len(self.val_loader.dataset)
        
        # Logging
        self.writer.add_scalar('val/loss', val_loss, epoch)
        self.writer.add_scalar('val/accuracy', accuracy, epoch)
        
        # Checkpointing
        if val_loss < self.best_val_loss:
            self.best_val_loss = val_loss
            self.save_checkpoint(epoch, is_best=True)
            
        return val_loss, accuracy
    
    def save_checkpoint(self, epoch, is_best=False):
        state = {
            'epoch': epoch,
            'model_state': self.model.state_dict(),
            'optimizer_state': self.optimizer.state_dict(),
            'scheduler_state': self.scheduler.state_dict(),
            'best_val_loss': self.best_val_loss,
            'config': self.config,
        }
        
        path = f"{self.config.checkpoint_dir}/checkpoint_epoch{epoch}.pt"
        torch.save(state, path)
        
        if is_best:
            best_path = f"{self.config.checkpoint_dir}/best_model.pt"
            torch.save(state, best_path)
    
    def run(self):
        for epoch in range(self.config.epochs):
            train_loss = self.train_epoch(epoch)
            val_loss, accuracy = self.validate(epoch)
            self.scheduler.step()
            
            print(f"Epoch {epoch}: train_loss={train_loss:.4f}, "
                  f"val_loss={val_loss:.4f}, accuracy={accuracy:.4f}")