# Step Function Contract

The composable function interfaces that `procedure.train_loop` and `procedure.train_epoch` / `val_epoch` accept. Understanding these contracts is essential for customizing training behavior.

## fn_train_step / fn_val_step

The most important customization point. Replaces the default `procedure.training_step` / `procedure.validation_step`.

### Signature

```python
def fn_train_step(
    ctx: FabricTrainingContext,
    batch: tuple[Tensor, ...],
    *,
    batch_idx: int,
    epoch_idx: int,
) -> tuple[Tensor | None, dict[str, float]]:
    ...
```

All parameters are passed as **keyword arguments** by `train_epoch`:
```python
y_pred, raw_metrics = fn_train_step(ctx=ctx, batch=batch, batch_idx=idx, epoch_idx=epoch_idx)
```

### Return value contract

| Index | Type | Description | Consumed by |
|-------|------|-------------|-------------|
| `[0]` | `Tensor \| None` | For normal supervised tasks, model predictions (`y_pred`). Shape must be compatible with torchmetrics in the recorder. If the recorder has no computed metrics and only tracks scalar losses, this can be a placeholder such as `None`. | `TorchMetricRecorder.update(pred=y_pred, target=y, **raw_metrics)` |
| `[1]` | `dict[str, float]` | Scalar metrics dict. Keys must match names defined in `scalar_metrics` of `TorchMetricRecorder`. | `TorchMetricRecorder.update(**raw_metrics)` via `ScalarMetric.update()` |

**Critical**: The dict keys in `[1]` must match the `scalar_metrics` configuration:
- Default recorder has `scalar_metrics='loss'` → step must return `{'loss': float}`
- If step returns `{'g_loss': ..., 'd_loss': ...}` → use `scalar_metrics=['g_loss', 'd_loss']` or `scalar_metrics='auto'`
- Mismatched keys are silently ignored (with a one-time warning), causing missing metrics in logs

For ordinary classification/regression tasks, the logic is:
- `fn_batch2target(batch)` extracts the target `y`
- `TorchMetricRecorder.update(y_pred, y, **raw_metrics)` uses `(y_pred, y)` for computed metrics such as Accuracy/F1
- `raw_metrics` is only for scalar logging such as `loss`, `ce_loss`, `reg_loss`

If `TorchMetricRecorder.metrics` is empty and the recorder is only tracking scalar losses, `pred`/`target` are not consumed for metric computation, so scalar-only procedures may return a placeholder `y_pred` such as `None`.

### Data flow diagram

```
train_epoch loop:
  for batch in train_loader:
    batch = fn_transform_batch(batch, batch_idx, epoch_idx)     # hook 1
    y_pred, raw_metrics = fn_train_step(ctx, batch, ...)        # hook 2
    y = fn_batch2target(batch)                                  # hook 3
    metrics = train_metric_recorder.update(y_pred, y, **raw_metrics)
    ctx.call_on_train_batch_end(batch, batch_idx, metrics)
  epoch_metrics = train_metric_recorder.compute()
```

### Default implementation (procedure.training_step)

Source: `common/procedure.py:83-128`

```python
def training_step(ctx, batch, *, fn_loss=None, batch_idx=None, **kwargs):
    ctx.model.train()
    x, y = batch
    loss_fn = _resolve_loss_fn(ctx, fn_loss)  # ctx.fn_loss or fn_loss param
    optimizer = ctx.optimizer
    optimizer.zero_grad()
    y_pred = ctx.model(x)
    loss = loss_fn(y_pred, y)
    ctx.fabric.backward(loss)
    optimizer.step()
    return y_pred, {'loss': loss.item()}
```

### Custom step example: multi-loss

```python
def my_train_step(ctx, batch, *, batch_idx, epoch_idx):
    x, y = batch
    ctx.model.train()
    ctx.optimizer.zero_grad()
    y_pred = ctx.model(x)
    ce_loss = F.cross_entropy(y_pred, y)
    reg_loss = 0.01 * sum(p.pow(2).sum() for p in ctx.model.parameters())
    total_loss = ce_loss + reg_loss
    ctx.fabric.backward(total_loss)
    ctx.optimizer.step()
    return y_pred, {'loss': total_loss.item(), 'ce_loss': ce_loss.item(), 'reg_loss': reg_loss.item()}
# Requires: scalar_metrics=['loss', 'ce_loss', 'reg_loss'] or scalar_metrics='auto'
```

### Validation step differences

`fn_val_step` has the same signature but:
- Should call `ctx.model.eval()` and use `torch.no_grad()`
- Should NOT call `optimizer.zero_grad()` / `optimizer.step()` / `fabric.backward()`
- Returns the same `tuple[Tensor, dict[str, float]]`

## fn_batch2target

Extracts the ground-truth target from a batch. Called after the step function to get `y` for metric computation.

### Signature

```python
def fn_batch2target(batch: tuple[Tensor, ...]) -> Tensor:
    ...
```

### Default

```python
def _default_batch_2_target(batch):
    return batch[1]  # assumes (x, y) tuple
```

### When to override

- **GAN training**: batch is `(real_data,)` with no labels. If the recorder only tracks scalar losses, the extracted target is effectively unused; existing examples commonly return `batch[0]` as a harmless placeholder.
- **Multi-input models**: batch is `(x1, x2, y)` → return `batch[2]`
- **Dict-based batches**: batch is `{'input': x, 'target': y}` → return `batch['target']`

### Wiring

```python
procedure.train_loop(ctx, train_loader, fn_batch2target=lambda batch: batch[2])
# Or for Trainer style:
class MyTrainer(BasicFabricTrainer):
    def batch_2_target(self, batch):
        return batch[2]
```

For scalar-only GAN-style training, the current examples use:

```python
procedure.train_loop(ctx, train_loader, fn_batch2target=lambda batch: batch[0])
```

That placeholder target is acceptable because the recorder is only tracking scalar losses and does not consume `(pred, target)` for torchmetrics.

## fn_transform_batch

Data augmentation hook called before the step function. Applied per-batch.

### Signature

```python
def fn_transform_batch(
    batch: tuple[Tensor, ...],
    batch_idx: int,
    epoch_idx: int | None,
) -> tuple[Tensor, ...]:
    ...
```

### Default

Identity function (no transformation). Two ways to provide:
1. Pass `fn_transform_batch` parameter to `train_epoch` / `val_epoch` (NOT available on `train_loop`)
2. Override `ctx.transform_on_train_batch_start` method — this is the recommended way when using `train_loop`, since `train_loop` delegates to `train_epoch` which reads this ctx method as fallback

### Example: Mixup

```python
def mixup_transform(batch, batch_idx, epoch_idx):
    x, y = batch
    lam = np.random.beta(0.2, 0.2)
    index = torch.randperm(x.size(0), device=x.device)
    x_mixed = lam * x + (1 - lam) * x[index]
    y_mixed = lam * y + (1 - lam) * y[index]
    return x_mixed, y_mixed

procedure.train_epoch(ctx, train_loader, fn_transform_batch=mixup_transform)
```

Note: For validation, there is a separate `ctx.transform_on_validation_batch_start` hook (e.g. for TTA).

## fn_on_trainepoch_end

Lightweight epoch-end hook. Simpler than writing a full callback.

### Signature

```python
def fn_on_trainepoch_end(
    ctx: FabricTrainingContext,
    epoch: int,
    train_metrics: dict[str, float],
    val_metrics: dict[str, float] | None,
) -> None:
    ...
```

### When to use

- Quick logging or printing custom info
- Adjusting hyperparameters mid-training
- For anything more complex (saving files, modifying training state), write a proper callback instead

### Wiring

```python
def my_epoch_end(ctx, epoch, train_metrics, val_metrics):
    if val_metrics and val_metrics.get('accuracy', 0) > 0.95:
        print(f'High accuracy reached at epoch {epoch}!')

procedure.train_loop(ctx, train_loader, fn_on_trainepoch_end=my_epoch_end)
```

## Composition summary

| Hook | Where injected | Default | Purpose |
|------|---------------|---------|---------|
| `fn_train_step` | `train_loop` / `train_epoch` | `procedure.training_step` | Core training logic |
| `fn_val_step` | `train_loop` / `val_epoch` | `procedure.validation_step` | Core validation logic |
| `fn_batch2target` | `train_loop` / `train_epoch` / `val_epoch` | `batch[1]` | Extract y from batch |
| `fn_transform_batch` | `train_epoch` / `val_epoch` | identity | Data augmentation |
| `fn_on_trainepoch_end` | `train_loop` only | None | Epoch-end custom logic |
