pipekit-train: support `eqx.nn.State` for stateful layers (BatchNorm, etc.) (v0.2, Q4)

## Problem / Request

Add support for `eqx.nn.State` so BatchNorm-like stateful layers can be trained. The v0.1 `TrainState` doesn't carry per-layer running statistics.

`boundaries.md` Q4: "BatchNorm and similar layers that carry running statistics need `eqx.nn.State`. The Equinox adapter's `TrainState` should probably carry an optional `eqx.nn.State` alongside the model."

## Motivation

- Many real models use BatchNorm (CNN classifiers, U-Nets). Without `eqx.nn.State` they can't be trained via this adapter.
- Straightforward extension; just adds a code path. The shape of the fix is documented.

## Proposed API

```python
class TrainState(eqx.Module):
model: eqx.Module
state: eqx.nn.State | None # NEW — optional stateful-layer state
opt_state: optax.OptState
step: jax.Array
```

`train_step` threads `state` through:

```python
def train_step(train_state, batch, key, task, optimizer):
(loss, (aux, new_state)) = grad_fn(train_state.model, batch, key, train_state.state)
...
return TrainState(model=new_model, state=new_state, ...), aux
```

`TrainTask.loss_fn` becomes optionally state-aware:
- If task has `stateful = True`: signature is `loss_fn(model, batch, key, state) -> (loss, (aux, new_state))`.
- Otherwise: existing v0.1 signature, ignored `state`.

`_SynthesisedTask` detects stateful layers via `eqx.is_inexact_array` partition and wires accordingly.

## Implementation Steps

- [ ] `TrainState.state: eqx.nn.State | None` field.
- [ ] `TrainState.create` builds initial state via `eqx.nn.State(model)` when the model contains stateful layers.
- [ ] `train_step` threads `state`; default `loss_fn` ignores it for backwards compat.
- [ ] `_SynthesisedTask` detects stateful models and adjusts the call.
- [ ] Save/restore preserves the `state` PyTree (extend Orbax bridge).
- [ ] Test: train a tiny CNN with `eqx.nn.BatchNorm`; assert running stats update across steps and are preserved through `save_state` / `restore_state`.

## Definition of Done

- [ ] BatchNorm-using model trains end-to-end via `TrainingLoop`.
- [ ] Existing stateless model paths unchanged.
- [ ] Docstring + ADR note.

## Relationships

- Parent: pipekit-train v0.2