From abacded5a442ad501cf193f8893f8236dff82e63 Mon Sep 17 00:00:00 2001
From: "Michael J. Williams" <michaeljw1@googlemail.com>
Date: Sat, 31 Oct 2026 09:49:16 +0100
Subject: [PATCH 1/4] ENH(zuko): restore optional flow training controls

Add gradient norm clipping, cosine learning-rate annealing, patience and opt-in progress, preserving best-epoch restoration and complete histories. Validate raw and transformed training particles before splitting.

Cover numerical optimization, RNG isolation and Sequince fit_kwargs integration. Document defaults and validation_fraction=0 behavior.

Validation: 42 Zuko tests and 93 full-suite tests pass on CPU; all pre-commit hooks and the strict Pixi docs build pass. GPU testing unavailable.

Refs: https://github.com/mj-will/aspire/issues/89
---
 CHANGELOG.md                       |   9 +
 docs/normalizing_flows.rst         |  58 ++++++
 pyproject.toml                     |   2 +
 src/sequince_density/flows/zuko.py | 162 ++++++++++++++---
 tests/test_zuko.py                 | 283 ++++++++++++++++++++++++++++-
 5 files changed, 484 insertions(+), 30 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index af05bb3..1899c75 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,15 @@ The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 ### Added
 
+- Optional Zuko training controls: total gradient-norm clipping (`clip_grad`),
+  cosine learning-rate annealing (`lr_annealing`), early stopping after
+  consecutive epochs without strict validation improvement (`patience`), and
+  opt-in epoch/loss display (`show_progress`). All are disabled by default and
+  work through Sequince flow adapters' `fit_kwargs`. Best validation parameters
+  are still restored and histories retain every completed epoch.
+- Zuko checks for invalid raw and transformed training data, with affected
+  dimensions reported for non-finite values before train/validation splitting.
+
 - Software archive DOI and method references for flow preconditioning and
   masked autoregressive flows.
 - Documentation sidebar links to the main Sequince documentation and
diff --git a/docs/normalizing_flows.rst b/docs/normalizing_flows.rst
index 94ba6ad..38fa7e2 100644
--- a/docs/normalizing_flows.rst
+++ b/docs/normalizing_flows.rst
@@ -34,6 +34,64 @@ Estimation*. Credit the architecture you select as well as the backend
 implementation: `Zuko <https://zuko.readthedocs.io/stable/>`_ or
 `FlowJAX <https://danielward27.github.io/flowjax/>`_.
 
+Torch training controls
+-----------------------
+
+``ZukoFlow.fit`` uses Adam with ``lr=1e-3``, ``n_epochs=100``,
+``batch_size=500`` and ``validation_fraction=0.2`` by default. Optional controls
+are disabled by default:
+
+* ``clip_grad=None``: set a finite positive value to clip the total gradient
+  norm across all parameters before each optimizer step.
+* ``lr_annealing=False``: set to ``True`` for cosine annealing from ``lr`` to
+  zero over ``n_epochs``. The scheduler steps after each epoch; the first epoch
+  uses ``lr``.
+* ``patience=None``: set a positive integer to stop after that many consecutive
+  epochs without a strictly lower validation loss. Ties count as unsuccessful
+  epochs and each improvement resets the counter.
+* ``show_progress=False``: set to ``True`` to display epoch progress and training
+  and validation losses. This uses the same option name as FlowJAX and does not
+  change seeded training results.
+
+For example:
+
+.. code-block:: python
+
+    flow.fit(
+        particles, n_epochs=100, clip_grad=5.0, lr_annealing=True,
+        patience=10, show_progress=True, rng=42,
+    )
+
+The parameters from the best validation epoch are restored even when early
+stopping is disabled. ``fit`` returns the adapter; ``flow.fit_history`` contains
+``training_loss`` and ``validation_loss`` lists for every completed epoch,
+including unsuccessful epochs before stopping. With ``validation_fraction=0``,
+all particles are used for both training and validation. Early stopping and
+best-epoch selection then use the loss evaluated on the training population
+after each epoch, rather than an independent holdout. The same fallback applies
+if rounding the split leaves either population empty.
+
+Training particles must have nonempty shape ``(n_samples, dims)``. Raw and
+transformed data are checked before splitting into training and validation
+populations. Non-finite values raise ``ValueError`` identifying the stage and
+affected dimensions (zero-based). Raw checks run before the particle transform
+can clip or obscure invalid values.
+
+All controls can also be passed in ``fit_kwargs`` to Sequince's
+``FlowTransform`` or ``FlowProposal``:
+
+.. code-block:: python
+
+    from sequince.transforms import FlowTransform
+
+    transform = FlowTransform(
+        ZukoFlow(dims=2, hidden_features=[32, 32]),
+        fit_kwargs={
+            "n_epochs": 100, "clip_grad": 5.0, "lr_annealing": True,
+            "patience": 10, "show_progress": True,
+        },
+    )
+
 Particle transform restrictions
 -------------------------------
 
diff --git a/pyproject.toml b/pyproject.toml
index 0b692df..3eb4b09 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -26,6 +26,7 @@ gmm = [
 zuko = [
     "torch>=2",
     "zuko>=1.6",
+    "tqdm>=4.66",
 ]
 flowjax = [
     "jax>0.5.3",
@@ -58,6 +59,7 @@ test = [
     "scikit-learn>=1.4",
     "torch>=2",
     "zuko>=1.6",
+    "tqdm>=4.66",
 ]
 dev = [
     "pre-commit",
diff --git a/src/sequince_density/flows/zuko.py b/src/sequince_density/flows/zuko.py
index f46bd9e..d0eb9c1 100644
--- a/src/sequince_density/flows/zuko.py
+++ b/src/sequince_density/flows/zuko.py
@@ -8,6 +8,7 @@
 import array_api_compat.torch as torch_api
 import torch
 import zuko
+from tqdm.auto import tqdm
 
 from ..utils import torch_random_context
 from .base import Flow
@@ -73,9 +74,53 @@ def fit(
         lr: float = 1e-3,
         batch_size: int = 500,
         validation_fraction: float = 0.2,
+        clip_grad: float | None = None,
+        lr_annealing: bool = False,
+        patience: int | None = None,
+        show_progress: bool = False,
         *,
         rng=None,
     ) -> "ZukoFlow":
+        """Fit by maximum likelihood and restore the best validation epoch.
+
+        Parameters
+        ----------
+        samples : array_like, shape (n_samples, dims)
+            Nonempty, finite training particles. Transformed particles must
+            also be finite and have the same shape.
+        context : object, optional
+            Context passed to the particle transform.
+        n_epochs : int, default=100
+            Maximum number of training epochs.
+        lr : float, default=1e-3
+            Initial Adam learning rate.
+        batch_size : int, default=500
+            Maximum number of particles per training or validation batch.
+        validation_fraction : float, default=0.2
+            Fraction held out for validation. If rounding gives no validation
+            particles (including zero), or leaves no training particles, use
+            all particles for both training and validation.
+        clip_grad : float or None, default=None
+            Positive maximum total gradient norm, applied before each optimizer
+            step. None disables clipping.
+        lr_annealing : bool, default=False
+            Apply cosine learning-rate annealing from lr to zero over n_epochs,
+            stepping the scheduler after each completed epoch.
+        patience : int or None, default=None
+            Positive number of consecutive epochs without a strictly lower
+            validation loss before stopping. Improvement resets the counter;
+            ties count as unsuccessful epochs. None disables early stopping.
+        show_progress : bool, default=False
+            Show epoch progress and training and validation losses.
+        rng : RandomGenerator or int or None, optional
+            RNG for fitting, isolated from global Torch state when supplied.
+
+        Returns
+        -------
+        ZukoFlow
+            This fitted adapter. fit_history contains losses for every
+            completed epoch, including epochs after the best validation epoch.
+        """
         with torch_random_context(rng, device=self.device):
             return self._fit(
                 samples,
@@ -84,6 +129,10 @@ def fit(
                 lr=lr,
                 batch_size=batch_size,
                 validation_fraction=validation_fraction,
+                clip_grad=clip_grad,
+                lr_annealing=lr_annealing,
+                patience=patience,
+                show_progress=show_progress,
                 rng=rng,
             )
 
@@ -95,10 +144,26 @@ def _fit(
         lr: float = 1e-3,
         batch_size: int = 500,
         validation_fraction: float = 0.2,
+        clip_grad: float | None = None,
+        lr_annealing: bool = False,
+        patience: int | None = None,
+        show_progress: bool = False,
         rng=None,
     ) -> "ZukoFlow":
+        if clip_grad is not None and (not 0 < clip_grad < float("inf")):
+            raise ValueError("clip_grad must be finite and positive or None.")
+        if patience is not None and (
+            isinstance(patience, bool)
+            or not isinstance(patience, int)
+            or patience < 1
+        ):
+            raise ValueError("patience must be a positive integer or None.")
         x = self._as_torch(samples)
+        self._validate_training_data(x, "Raw")
         x_prime = self.fit_particle_transform(x, context=context, rng=rng)
+        self._validate_training_data(x_prime, "Transformed")
+        if x_prime.shape != x.shape:
+            raise ValueError("Transformed training data must preserve shape.")
         indices = torch.randperm(x_prime.shape[0], device=x_prime.device)
         x_prime = x_prime[indices]
 
@@ -122,40 +187,65 @@ def _fit(
         )
 
         optimizer = torch.optim.Adam(self.flow.parameters(), lr=lr)
+        scheduler = (
+            torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, n_epochs)
+            if lr_annealing
+            else None
+        )
+        unsuccessful_epochs = 0
         best_state = None
         best_val_loss = float("inf")
         training_loss = []
         validation_loss = []
 
-        for _ in range(n_epochs):
-            self.flow.train()
-            train_loss = 0.0
-            n_train_batches = 0
-            for (x_batch,) in dataset:
-                loss = self.loss_fn(x_batch)
-                optimizer.zero_grad()
-                loss.backward()
-                optimizer.step()
-                train_loss += loss.item()
-                n_train_batches += 1
-
-            self.flow.eval()
-            val_loss = 0.0
-            n_batches = 0
-            for (x_batch,) in val_dataset:
-                with torch.no_grad():
-                    val_loss += self.loss_fn(x_batch).item()
-                n_batches += 1
-
-            avg_val_loss = val_loss / max(n_batches, 1)
-            training_loss.append(train_loss / max(n_train_batches, 1))
-            validation_loss.append(avg_val_loss)
-            if avg_val_loss < best_val_loss:
-                best_val_loss = avg_val_loss
-                best_state = {
-                    key: value.detach().clone()
-                    for key, value in self.flow.state_dict().items()
-                }
+        with tqdm(
+            range(n_epochs), desc="Epochs", disable=not show_progress
+        ) as progress:
+            for _ in progress:
+                self.flow.train()
+                train_loss = 0.0
+                n_train_batches = 0
+                for (x_batch,) in dataset:
+                    loss = self.loss_fn(x_batch)
+                    optimizer.zero_grad()
+                    loss.backward()
+                    if clip_grad is not None:
+                        torch.nn.utils.clip_grad_norm_(
+                            self.flow.parameters(), clip_grad
+                        )
+                    optimizer.step()
+                    train_loss += loss.item()
+                    n_train_batches += 1
+
+                if scheduler is not None:
+                    scheduler.step()
+                self.flow.eval()
+                val_loss = 0.0
+                n_batches = 0
+                for (x_batch,) in val_dataset:
+                    with torch.no_grad():
+                        val_loss += self.loss_fn(x_batch).item()
+                    n_batches += 1
+
+                avg_val_loss = val_loss / max(n_batches, 1)
+                training_loss.append(train_loss / max(n_train_batches, 1))
+                validation_loss.append(avg_val_loss)
+                if avg_val_loss < best_val_loss:
+                    unsuccessful_epochs = 0
+                    best_val_loss = avg_val_loss
+                    best_state = {
+                        key: value.detach().clone()
+                        for key, value in self.flow.state_dict().items()
+                    }
+                else:
+                    unsuccessful_epochs += 1
+                if show_progress:
+                    progress.set_postfix(
+                        train_loss=f"{training_loss[-1]:.4f}",
+                        val_loss=f"{avg_val_loss:.4f}",
+                    )
+                if patience is not None and unsuccessful_epochs >= patience:
+                    break
 
         if best_state is not None:
             self.flow.load_state_dict(best_state)
@@ -167,6 +257,20 @@ def _fit(
         self.flow.eval()
         return self
 
+    def _validate_training_data(self, x: torch.Tensor, stage: str) -> None:
+        if x.ndim != 2 or x.shape[0] == 0 or x.shape[1] != self.dims:
+            raise ValueError(
+                f"{stage} training data must have nonempty shape "
+                f"(n_samples, {self.dims})."
+            )
+        invalid = ~torch.isfinite(x)
+        if invalid.any():
+            dimensions = invalid.any(dim=0).nonzero(as_tuple=True)[0].tolist()
+            raise ValueError(
+                f"{stage} training data contains non-finite values "
+                f"in dimensions: {dimensions}"
+            )
+
     def state_dict(self):
         return {
             "flow": {
diff --git a/tests/test_zuko.py b/tests/test_zuko.py
index 0f3662c..30b96cc 100644
--- a/tests/test_zuko.py
+++ b/tests/test_zuko.py
@@ -2,10 +2,12 @@
 
 import pickle
 from pathlib import Path
+from typing import Any
 
 import numpy as np
 import pytest
 from orng import RandomGenerator
+from sequince.proposals import FlowProposal
 from sequince.transforms import (
     AffineTransform,
     CompositeTransform,
@@ -13,9 +15,11 @@
     ParticleTransform,
 )
 
-torch = pytest.importorskip("torch")
+pytest.importorskip("torch")
 pytest.importorskip("zuko")
 
+import torch  # noqa: E402
+
 from sequince_density.flows.zuko import ZukoFlow  # noqa: E402
 
 
@@ -143,3 +147,280 @@ def test_checkpoint_rejects_periodic_particle_transform(
     torch.testing.assert_close(
         flow.log_prob(points), original_log_prob, rtol=0, atol=0
     )
+
+
+class TrainableGaussian(torch.nn.Module):
+    """Unit-variance Gaussian with an observable, trainable location."""
+
+    def __init__(self, dims: int, context: int = 0) -> None:
+        super().__init__()
+        # Separate parameters distinguish total-norm from per-parameter clipping.
+        self.coordinates = torch.nn.ParameterList(
+            [torch.nn.Parameter(torch.zeros(())) for _ in range(dims)]
+        )
+        self.validation_locations: list[torch.Tensor] = []
+        self.gradient_norms: list[float] = []
+
+    @property
+    def loc(self) -> torch.Tensor:
+        return torch.stack(list(self.coordinates))
+
+    def forward(self) -> torch.distributions.Distribution:
+        if not self.training:
+            self.validation_locations.append(self.loc.detach().clone())
+            gradients = torch.stack([p.grad for p in self.coordinates])
+            self.gradient_norms.append(float(gradients.norm()))
+        return torch.distributions.Independent(
+            torch.distributions.Normal(self.loc, torch.ones_like(self.loc)), 1
+        )
+
+
+def gaussian_flow() -> ZukoFlow:
+    return ZukoFlow(
+        dims=2, flow_class=TrainableGaussian, dtype=torch.float64, device="cpu"
+    )
+
+
+@pytest.mark.parametrize("clip_grad", [None, 0.25])
+def test_gradient_clipping_bounds_total_norm(clip_grad: float | None) -> None:
+    flow = gaussian_flow()
+    flow.fit(
+        torch.tensor([[5.0, -12.0]]).repeat(8, 1),
+        n_epochs=4,
+        lr=0.1,
+        batch_size=8,
+        validation_fraction=0,
+        clip_grad=clip_grad,
+        rng=42,
+    )
+    norms = flow.flow.gradient_norms
+    assert len(norms) == 4
+    if clip_grad is None:
+        assert min(norms) > 10
+    else:
+        assert all(0 < norm <= clip_grad for norm in norms)
+    assert not torch.equal(flow.flow.loc, torch.zeros_like(flow.flow.loc))
+
+
+@pytest.mark.parametrize("lr_annealing", [False, True])
+def test_annealing_matches_adam_updates(lr_annealing: bool) -> None:
+    # With one batch per epoch, the Gaussian's gradient is loc - target.
+    target = torch.tensor([10.0, -10.0], dtype=torch.float64)
+    flow = gaussian_flow()
+    flow.fit(
+        target.repeat(8, 1),
+        n_epochs=5,
+        lr=0.1,
+        batch_size=8,
+        validation_fraction=0,
+        lr_annealing=lr_annealing,
+        rng=42,
+    )
+    loc = torch.zeros_like(target)
+    first = torch.zeros_like(target)
+    second = torch.zeros_like(target)
+    for epoch, actual in enumerate(flow.flow.validation_locations):
+        grad = loc - target
+        first = 0.9 * first + 0.1 * grad
+        second = 0.999 * second + 0.001 * grad.square()
+        rate = 0.1
+        if lr_annealing:
+            rate *= (1 + np.cos(np.pi * epoch / 5)) / 2
+        loc -= (
+            rate
+            * (first / (1 - 0.9 ** (epoch + 1)))
+            / ((second / (1 - 0.999 ** (epoch + 1))).sqrt() + 1e-8)
+        )
+        torch.testing.assert_close(actual, loc, rtol=1e-12, atol=1e-12)
+    torch.testing.assert_close(flow.flow.loc, loc, rtol=0, atol=1e-12)
+
+
+@pytest.mark.parametrize("patience", [None, 1, 3])
+def test_patience_counts_ties(patience: int | None) -> None:
+    flow = gaussian_flow()
+    flow.fit(
+        torch.ones(8, 2),
+        n_epochs=8,
+        lr=0,
+        validation_fraction=0,
+        patience=patience,
+        rng=42,
+    )
+    expected = 8 if patience is None else 1 + patience
+    assert len(flow.fit_history["training_loss"]) == expected
+    assert len(flow.fit_history["validation_loss"]) == expected
+
+
+@pytest.mark.parametrize("patience", [None, 1])
+def test_patience_allows_continued_improvement(patience: int | None) -> None:
+    flow = gaussian_flow()
+    flow.fit(
+        torch.full((8, 2), 10.0),
+        n_epochs=8,
+        lr=0.1,
+        validation_fraction=0,
+        patience=patience,
+        rng=42,
+    )
+    losses = flow.fit_history["validation_loss"]
+    assert len(losses) == 8
+    assert all(a > b for a, b in zip(losses, losses[1:]))
+
+
+@pytest.mark.parametrize("patience", [None, 3, 5])
+def test_restores_best_validation_parameters(patience: int | None) -> None:
+    flow = gaussian_flow()
+    flow.fit(
+        torch.full((8, 2), 0.25),
+        n_epochs=20,
+        lr=0.2,
+        validation_fraction=0,
+        patience=patience,
+        rng=42,
+    )
+    losses = flow.fit_history["validation_loss"]
+    locations = flow.flow.validation_locations
+    best = int(np.argmin(losses))
+    assert best < len(losses) - 1
+    torch.testing.assert_close(flow.flow.loc, locations[best], rtol=0, atol=0)
+    assert len(flow.fit_history["training_loss"]) == len(losses)
+    if patience == 5:
+        # Four unsuccessful epochs precede a new best at epoch six.
+        # The counter must reset there, allowing five more epochs.
+        assert best == 5
+        assert len(losses) == 11
+    if patience is not None:
+        assert len(losses) == best + 1 + patience
+    else:
+        assert len(losses) == 20
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+def test_progress_preserves_seeded_real_zuko_training(
+    dtype: torch.dtype, capsys: pytest.CaptureFixture[str]
+) -> None:
+    particles = np.random.default_rng(21).normal(size=(32, 2))
+    flows = []
+    for show_progress in (False, True):
+        flow = ZukoFlow(dims=2, hidden_features=[8, 8], dtype=dtype)
+        state = torch.random.get_rng_state().clone()
+        assert (
+            flow.fit(
+                particles,
+                n_epochs=3,
+                batch_size=8,
+                clip_grad=0.5,
+                lr_annealing=True,
+                patience=2,
+                show_progress=show_progress,
+                rng=42,
+            )
+            is flow
+        )
+        assert torch.equal(torch.random.get_rng_state(), state)
+        output = capsys.readouterr().err
+        if show_progress:
+            assert "Epochs" in output
+            assert "train_loss=" in output and "val_loss=" in output
+        else:
+            assert output == ""
+        assert all(
+            p.dtype == dtype and p.device.type == "cpu"
+            for p in flow.flow.parameters()
+        )
+        flows.append(flow)
+    assert flows[0].fit_history == flows[1].fit_history
+    for key, value in flows[0].flow.state_dict().items():
+        torch.testing.assert_close(
+            value, flows[1].flow.state_dict()[key], rtol=0, atol=0
+        )
+
+
+@pytest.mark.parametrize("adapter", [FlowTransform, FlowProposal])
+def test_training_controls_through_sequince(adapter: Any) -> None:
+    particles = np.random.default_rng(21).normal(size=(32, 2))
+    options: dict[str, Any] = dict(
+        n_epochs=3,
+        batch_size=8,
+        clip_grad=0.5,
+        lr_annealing=True,
+        patience=2,
+        show_progress=False,
+    )
+    direct = ZukoFlow(dims=2, hidden_features=[8, 8], seed=22)
+    direct.fit(particles, rng=42, **options)
+    wrapped = adapter(
+        ZukoFlow(dims=2, hidden_features=[8, 8], seed=22), fit_kwargs=options
+    ).fit(particles, rng=42)
+    assert wrapped.flow.fit_history == direct.fit_history
+    for key, value in direct.flow.state_dict().items():
+        torch.testing.assert_close(
+            value, wrapped.flow.flow.state_dict()[key], rtol=0, atol=0
+        )
+
+
+@pytest.mark.parametrize(
+    "bad_value", [float("nan"), float("inf"), -float("inf")]
+)
+def test_rejects_nonfinite_raw_data(bad_value: float) -> None:
+    # The bounded logit clips infinities, so validation must precede it.
+    from sequince.transforms import BoundedToUnboundedTransform
+
+    flow = ZukoFlow(
+        dims=3,
+        hidden_features=[8, 8],
+        particle_transform=BoundedToUnboundedTransform(
+            {i: (0, 1) for i in range(3)}
+        ),
+    )
+    particles = torch.full((8, 3), 0.5)
+    particles[0, 0] = bad_value
+    particles[-1, 2] = bad_value
+    with pytest.raises(
+        ValueError, match=r"Raw.*non-finite.*dimensions: \[0, 2\]"
+    ):
+        flow.fit(particles, n_epochs=1, rng=42)
+
+
+class NonfiniteTransform(ParticleTransform):
+    def forward(
+        self, particles: torch.Tensor, context: Any = None, *, rng: Any = None
+    ) -> tuple[torch.Tensor, torch.Tensor, None]:
+        values = particles.clone()
+        values[0, 0] = float("nan")
+        values[-1, 2] = float("inf")
+        return values, torch.zeros(len(values)), None
+
+
+def test_rejects_nonfinite_transformed_data() -> None:
+    flow = ZukoFlow(dims=3, particle_transform=NonfiniteTransform())
+    with pytest.raises(
+        ValueError, match=r"Transformed.*non-finite.*dimensions: \[0, 2\]"
+    ):
+        flow.fit(torch.zeros(8, 3), n_epochs=1, rng=42)
+
+
+@pytest.mark.parametrize(
+    "samples", [[], [1, 2], np.zeros((0, 2)), np.zeros((8, 3))]
+)
+def test_rejects_invalid_training_shapes(samples: Any) -> None:
+    with pytest.raises(ValueError, match="Raw training data.*shape"):
+        gaussian_flow().fit(samples, n_epochs=1)
+
+
+@pytest.mark.parametrize(
+    "options",
+    [
+        {"clip_grad": 0},
+        {"clip_grad": -1},
+        {"clip_grad": float("nan")},
+        {"clip_grad": float("inf")},
+        {"patience": 0},
+        {"patience": -1},
+        {"patience": 1.5},
+        {"patience": True},
+    ],
+)
+def test_rejects_invalid_training_controls(options: dict[str, Any]) -> None:
+    with pytest.raises(ValueError, match="clip_grad|patience"):
+        gaussian_flow().fit(torch.zeros(8, 2), n_epochs=1, **options)

From 2bacca70ff1dd07669cd2b5f09fc12cdd42ea3bd Mon Sep 17 00:00:00 2001
From: "Michael J. Williams" <michaeljw1@googlemail.com>
Date: Sat, 31 Oct 2026 09:51:11 +0100
Subject: [PATCH 2/4] MAINT(flows): share backend-neutral training data
 validation

Extract the Zuko validation helper into density utilities using the native array namespace. Preserve shape checks and dimension-specific errors, and cover NumPy, Torch and JAX inputs.

Validation: 96 full-suite tests passed; all pre-commit checks and the strict Pixi docs build passed. GPU testing unavailable.
---
 src/sequince_density/flows/zuko.py | 20 ++-----------
 src/sequince_density/utils.py      | 46 ++++++++++++++++++++++++++++--
 tests/test_training_data.py        | 28 ++++++++++++++++++
 3 files changed, 75 insertions(+), 19 deletions(-)
 create mode 100644 tests/test_training_data.py

diff --git a/src/sequince_density/flows/zuko.py b/src/sequince_density/flows/zuko.py
index d0eb9c1..0d00766 100644
--- a/src/sequince_density/flows/zuko.py
+++ b/src/sequince_density/flows/zuko.py
@@ -10,7 +10,7 @@
 import zuko
 from tqdm.auto import tqdm
 
-from ..utils import torch_random_context
+from ..utils import torch_random_context, validate_training_data
 from .base import Flow
 
 logger = logging.getLogger(__name__)
@@ -159,9 +159,9 @@ def _fit(
         ):
             raise ValueError("patience must be a positive integer or None.")
         x = self._as_torch(samples)
-        self._validate_training_data(x, "Raw")
+        validate_training_data(x, self.dims, "Raw")
         x_prime = self.fit_particle_transform(x, context=context, rng=rng)
-        self._validate_training_data(x_prime, "Transformed")
+        validate_training_data(x_prime, self.dims, "Transformed")
         if x_prime.shape != x.shape:
             raise ValueError("Transformed training data must preserve shape.")
         indices = torch.randperm(x_prime.shape[0], device=x_prime.device)
@@ -257,20 +257,6 @@ def _fit(
         self.flow.eval()
         return self
 
-    def _validate_training_data(self, x: torch.Tensor, stage: str) -> None:
-        if x.ndim != 2 or x.shape[0] == 0 or x.shape[1] != self.dims:
-            raise ValueError(
-                f"{stage} training data must have nonempty shape "
-                f"(n_samples, {self.dims})."
-            )
-        invalid = ~torch.isfinite(x)
-        if invalid.any():
-            dimensions = invalid.any(dim=0).nonzero(as_tuple=True)[0].tolist()
-            raise ValueError(
-                f"{stage} training data contains non-finite values "
-                f"in dimensions: {dimensions}"
-            )
-
     def state_dict(self):
         return {
             "flow": {
diff --git a/src/sequince_density/utils.py b/src/sequince_density/utils.py
index a321d64..a8c3fd9 100644
--- a/src/sequince_density/utils.py
+++ b/src/sequince_density/utils.py
@@ -7,7 +7,49 @@
 from typing import Any
 
 from orng import RandomGenerator
-from sequince.array_api import to_numpy
+from sequince.array_api import namespace, to_numpy
+from sequince.typing import Array
+
+
+def validate_training_data(
+    samples: Array, dims: int, stage: str = "Raw"
+) -> None:
+    """Validate training particles using their native array backend.
+
+    Parameters
+    ----------
+    samples : array, shape (n_samples, dims)
+        Nonempty, finite particle coordinates. Validation does not modify them.
+    dims : int
+        Expected number of dimensions.
+    stage : str, default="Raw"
+        Label for the data in error messages, for example "Transformed".
+
+    Raises
+    ------
+    ValueError
+        If the shape is invalid or any coordinates are non-finite. Non-finite
+        errors report the affected dimension indices (zero-based).
+    """
+    if samples.ndim != 2 or samples.shape[0] == 0 or samples.shape[1] != dims:
+        raise ValueError(
+            f"{stage} training data must have nonempty shape "
+            f"(n_samples, {dims})."
+        )
+    xp = namespace(samples)
+    invalid = ~xp.isfinite(samples)
+    if xp.any(invalid):
+        dimensions = [
+            i
+            for i, affected in enumerate(
+                to_numpy(xp.any(invalid, axis=0)).tolist()
+            )
+            if affected
+        ]
+        raise ValueError(
+            f"{stage} training data contains non-finite values "
+            f"in dimensions: {dimensions}"
+        )
 
 
 def draw_seed(rng: RandomGenerator | int | None) -> int | None:
@@ -47,4 +89,4 @@ def torch_random_context(
         yield
 
 
-__all__ = ["draw_seed", "torch_random_context"]
+__all__ = ["draw_seed", "torch_random_context", "validate_training_data"]
diff --git a/tests/test_training_data.py b/tests/test_training_data.py
new file mode 100644
index 0000000..b47dae7
--- /dev/null
+++ b/tests/test_training_data.py
@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from sequince_density.utils import validate_training_data
+
+
+@pytest.mark.parametrize("backend", ["numpy", "torch", "jax.numpy"])
+def test_validate_training_data_across_backends(backend: str) -> None:
+    xp = pytest.importorskip(backend)
+    values = np.array([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]])
+    particles = xp.asarray(values)
+    validate_training_data(particles, dims=3)
+    np.testing.assert_array_equal(np.asarray(particles), values)
+
+    invalid = values.copy()
+    invalid[0, 0] = np.nan
+    invalid[1, 2] = np.inf
+    with pytest.raises(
+        ValueError, match=r"Transformed.*non-finite.*dimensions: \[0, 2\]"
+    ):
+        validate_training_data(
+            xp.asarray(invalid), dims=3, stage="Transformed"
+        )
+
+    with pytest.raises(ValueError, match="Raw training data.*shape"):
+        validate_training_data(particles, dims=2)

From e5f3131ffe7c20f96a373d1d5a8ce94ae26fb8eb Mon Sep 17 00:00:00 2001
From: "Michael J. Williams" <michaeljw1@googlemail.com>
Date: Sat, 31 Oct 2026 10:15:57 +0100
Subject: [PATCH 3/4] DOC: shorten training control documentation

---
 CHANGELOG.md               | 13 +++-----
 docs/normalizing_flows.rst | 64 ++++++++++++--------------------------
 2 files changed, 24 insertions(+), 53 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1899c75..799be94 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,15 +8,10 @@ The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 ### Added
 
-- Optional Zuko training controls: total gradient-norm clipping (`clip_grad`),
-  cosine learning-rate annealing (`lr_annealing`), early stopping after
-  consecutive epochs without strict validation improvement (`patience`), and
-  opt-in epoch/loss display (`show_progress`). All are disabled by default and
-  work through Sequince flow adapters' `fit_kwargs`. Best validation parameters
-  are still restored and histories retain every completed epoch.
-- Zuko checks for invalid raw and transformed training data, with affected
-  dimensions reported for non-finite values before train/validation splitting.
-
+- Optional Zuko gradient clipping, cosine learning-rate annealing, early stopping
+  and progress display, available through flow adapters' `fit_kwargs`.
+- Zuko training-data checks before and after particle transforms, reporting
+  affected dimensions for non-finite values.
 - Software archive DOI and method references for flow preconditioning and
   masked autoregressive flows.
 - Documentation sidebar links to the main Sequince documentation and
diff --git a/docs/normalizing_flows.rst b/docs/normalizing_flows.rst
index 38fa7e2..51d0bf3 100644
--- a/docs/normalizing_flows.rst
+++ b/docs/normalizing_flows.rst
@@ -37,23 +37,15 @@ implementation: `Zuko <https://zuko.readthedocs.io/stable/>`_ or
 Torch training controls
 -----------------------
 
-``ZukoFlow.fit`` uses Adam with ``lr=1e-3``, ``n_epochs=100``,
-``batch_size=500`` and ``validation_fraction=0.2`` by default. Optional controls
-are disabled by default:
-
-* ``clip_grad=None``: set a finite positive value to clip the total gradient
-  norm across all parameters before each optimizer step.
-* ``lr_annealing=False``: set to ``True`` for cosine annealing from ``lr`` to
-  zero over ``n_epochs``. The scheduler steps after each epoch; the first epoch
-  uses ``lr``.
-* ``patience=None``: set a positive integer to stop after that many consecutive
-  epochs without a strictly lower validation loss. Ties count as unsuccessful
-  epochs and each improvement resets the counter.
-* ``show_progress=False``: set to ``True`` to display epoch progress and training
-  and validation losses. This uses the same option name as FlowJAX and does not
-  change seeded training results.
-
-For example:
+``ZukoFlow.fit`` accepts these optional controls, all disabled by default:
+
+* ``clip_grad=None``: set a finite positive total gradient-norm limit.
+* ``lr_annealing=False``: enable cosine annealing over ``n_epochs``, starting at
+  ``lr`` and stepping after each epoch.
+* ``patience=None``: set a positive number of consecutive epochs without a
+  strictly lower validation loss before stopping. Ties count; improvement
+  resets the counter.
+* ``show_progress=False``: display epochs and training/validation losses.
 
 .. code-block:: python
 
@@ -62,35 +54,19 @@ For example:
         patience=10, show_progress=True, rng=42,
     )
 
-The parameters from the best validation epoch are restored even when early
-stopping is disabled. ``fit`` returns the adapter; ``flow.fit_history`` contains
-``training_loss`` and ``validation_loss`` lists for every completed epoch,
-including unsuccessful epochs before stopping. With ``validation_fraction=0``,
-all particles are used for both training and validation. Early stopping and
-best-epoch selection then use the loss evaluated on the training population
-after each epoch, rather than an independent holdout. The same fallback applies
-if rounding the split leaves either population empty.
-
-Training particles must have nonempty shape ``(n_samples, dims)``. Raw and
-transformed data are checked before splitting into training and validation
-populations. Non-finite values raise ``ValueError`` identifying the stage and
-affected dimensions (zero-based). Raw checks run before the particle transform
-can clip or obscure invalid values.
-
-All controls can also be passed in ``fit_kwargs`` to Sequince's
-``FlowTransform`` or ``FlowProposal``:
-
-.. code-block:: python
+Pass the same options through ``fit_kwargs`` in ``FlowTransform`` or
+``FlowProposal``. See :py:meth:`~sequince_density.flows.zuko.ZukoFlow.fit`
+for all training defaults.
 
-    from sequince.transforms import FlowTransform
+The best validation parameters are always restored. ``fit`` returns the adapter,
+with losses for every completed epoch in ``fit_history``. With
+``validation_fraction=0`` (or an empty split after rounding), training and
+validation both use all particles, including for early stopping and model
+selection.
 
-    transform = FlowTransform(
-        ZukoFlow(dims=2, hidden_features=[32, 32]),
-        fit_kwargs={
-            "n_epochs": 100, "clip_grad": 5.0, "lr_annealing": True,
-            "patience": 10, "show_progress": True,
-        },
-    )
+Raw and transformed particles must be finite with nonempty shape
+``(n_samples, dims)``. Invalid data raises ``ValueError``; non-finite errors
+identify the stage and affected dimensions (zero-based).
 
 Particle transform restrictions
 -------------------------------

From c71791660e4d637b01eaa79b59c7061e4c5bd9b2 Mon Sep 17 00:00:00 2001
From: "Michael J. Williams" <michaeljw1@googlemail.com>
Date: Sat, 31 Oct 2026 10:25:23 +0100
Subject: [PATCH 4/4] BUG(zuko): skip validation when its fraction is zero

Keep final parameters and an empty validation history when validation is disabled. Reject patience in that mode; retain best-epoch selection when validation is enabled.

Validation: 50 Zuko tests and 104 full-suite tests passed on CPU, with no skips. All pre-commit hooks and the strict docs build passed. GPU testing unavailable.
---
 CHANGELOG.md                       |  6 +++
 docs/normalizing_flows.rst         | 11 +++--
 src/sequince_density/flows/zuko.py | 41 +++++++++++-----
 tests/test_zuko.py                 | 77 ++++++++++++++++++++++++++----
 4 files changed, 108 insertions(+), 27 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 799be94..e29ba76 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -17,6 +17,12 @@ The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 - Documentation sidebar links to the main Sequince documentation and
   sequince-examples under Related Packages.
 
+### Changed
+
+- Zuko's `validation_fraction=0` skips validation and keeps final parameters,
+  with an empty validation-loss history. Combining it with `patience` raises
+  `ValueError`.
+
 ## [0.1.0b1] - 2026-10-27
 
 First public beta release.
diff --git a/docs/normalizing_flows.rst b/docs/normalizing_flows.rst
index 51d0bf3..fc23565 100644
--- a/docs/normalizing_flows.rst
+++ b/docs/normalizing_flows.rst
@@ -58,11 +58,12 @@ Pass the same options through ``fit_kwargs`` in ``FlowTransform`` or
 ``FlowProposal``. See :py:meth:`~sequince_density.flows.zuko.ZukoFlow.fit`
 for all training defaults.
 
-The best validation parameters are always restored. ``fit`` returns the adapter,
-with losses for every completed epoch in ``fit_history``. With
-``validation_fraction=0`` (or an empty split after rounding), training and
-validation both use all particles, including for early stopping and model
-selection.
+``fit`` returns the adapter with losses for every completed epoch in
+``fit_history`` and restores the best validation parameters. With
+``validation_fraction=0``, validation is skipped, ``validation_loss`` is empty,
+and final parameters are retained; setting ``patience`` raises ``ValueError``.
+For nonzero fractions, an empty split after rounding uses all particles for
+both training and validation.
 
 Raw and transformed particles must be finite with nonempty shape
 ``(n_samples, dims)``. Invalid data raises ``ValueError``; non-finite errors
diff --git a/src/sequince_density/flows/zuko.py b/src/sequince_density/flows/zuko.py
index 0d00766..9840f79 100644
--- a/src/sequince_density/flows/zuko.py
+++ b/src/sequince_density/flows/zuko.py
@@ -81,7 +81,7 @@ def fit(
         *,
         rng=None,
     ) -> "ZukoFlow":
-        """Fit by maximum likelihood and restore the best validation epoch.
+        """Fit by maximum likelihood, restoring the best epoch when validating.
 
         Parameters
         ----------
@@ -97,9 +97,9 @@ def fit(
         batch_size : int, default=500
             Maximum number of particles per training or validation batch.
         validation_fraction : float, default=0.2
-            Fraction held out for validation. If rounding gives no validation
-            particles (including zero), or leaves no training particles, use
-            all particles for both training and validation.
+            Fraction held out for validation. Zero skips validation and retains
+            the final parameters. For nonzero fractions, an empty split after
+            rounding uses all particles for both training and validation.
         clip_grad : float or None, default=None
             Positive maximum total gradient norm, applied before each optimizer
             step. None disables clipping.
@@ -109,7 +109,8 @@ def fit(
         patience : int or None, default=None
             Positive number of consecutive epochs without a strictly lower
             validation loss before stopping. Improvement resets the counter;
-            ties count as unsuccessful epochs. None disables early stopping.
+            ties count as unsuccessful epochs. Requires validation; None
+            disables early stopping.
         show_progress : bool, default=False
             Show epoch progress and training and validation losses.
         rng : RandomGenerator or int or None, optional
@@ -119,7 +120,7 @@ def fit(
         -------
         ZukoFlow
             This fitted adapter. fit_history contains losses for every
-            completed epoch, including epochs after the best validation epoch.
+            completed epoch; validation_loss is empty when validation is off.
         """
         with torch_random_context(rng, device=self.device):
             return self._fit(
@@ -158,6 +159,8 @@ def _fit(
             or patience < 1
         ):
             raise ValueError("patience must be a positive integer or None.")
+        if validation_fraction == 0 and patience is not None:
+            raise ValueError("patience requires validation_fraction > 0.")
         x = self._as_torch(samples)
         validate_training_data(x, self.dims, "Raw")
         x_prime = self.fit_particle_transform(x, context=context, rng=rng)
@@ -168,7 +171,10 @@ def _fit(
         x_prime = x_prime[indices]
 
         n_val = int(validation_fraction * x_prime.shape[0])
-        if n_val <= 0 or n_val >= x_prime.shape[0]:
+        if validation_fraction == 0:
+            x_train = x_prime
+            x_val = None
+        elif n_val <= 0 or n_val >= x_prime.shape[0]:
             x_train = x_prime
             x_val = x_prime
         else:
@@ -180,10 +186,14 @@ def _fit(
             shuffle=True,
             batch_size=min(batch_size, x_train.shape[0]),
         )
-        val_dataset = torch.utils.data.DataLoader(
-            torch.utils.data.TensorDataset(x_val),
-            shuffle=False,
-            batch_size=min(batch_size, x_val.shape[0]),
+        val_dataset = (
+            torch.utils.data.DataLoader(
+                torch.utils.data.TensorDataset(x_val),
+                shuffle=False,
+                batch_size=min(batch_size, x_val.shape[0]),
+            )
+            if x_val is not None
+            else None
         )
 
         optimizer = torch.optim.Adam(self.flow.parameters(), lr=lr)
@@ -219,6 +229,14 @@ def _fit(
 
                 if scheduler is not None:
                     scheduler.step()
+                training_loss.append(train_loss / max(n_train_batches, 1))
+                if val_dataset is None:
+                    if show_progress:
+                        progress.set_postfix(
+                            train_loss=f"{training_loss[-1]:.4f}"
+                        )
+                    continue
+
                 self.flow.eval()
                 val_loss = 0.0
                 n_batches = 0
@@ -228,7 +246,6 @@ def _fit(
                     n_batches += 1
 
                 avg_val_loss = val_loss / max(n_batches, 1)
-                training_loss.append(train_loss / max(n_train_batches, 1))
                 validation_loss.append(avg_val_loss)
                 if avg_val_loss < best_val_loss:
                     unsuccessful_epochs = 0
diff --git a/tests/test_zuko.py b/tests/test_zuko.py
index 30b96cc..e6b9b0b 100644
--- a/tests/test_zuko.py
+++ b/tests/test_zuko.py
@@ -189,7 +189,7 @@ def test_gradient_clipping_bounds_total_norm(clip_grad: float | None) -> None:
         n_epochs=4,
         lr=0.1,
         batch_size=8,
-        validation_fraction=0,
+        validation_fraction=0.25,
         clip_grad=clip_grad,
         rng=42,
     )
@@ -212,7 +212,7 @@ def test_annealing_matches_adam_updates(lr_annealing: bool) -> None:
         n_epochs=5,
         lr=0.1,
         batch_size=8,
-        validation_fraction=0,
+        validation_fraction=0.25,
         lr_annealing=lr_annealing,
         rng=42,
     )
@@ -242,7 +242,7 @@ def test_patience_counts_ties(patience: int | None) -> None:
         torch.ones(8, 2),
         n_epochs=8,
         lr=0,
-        validation_fraction=0,
+        validation_fraction=0.25,
         patience=patience,
         rng=42,
     )
@@ -258,7 +258,7 @@ def test_patience_allows_continued_improvement(patience: int | None) -> None:
         torch.full((8, 2), 10.0),
         n_epochs=8,
         lr=0.1,
-        validation_fraction=0,
+        validation_fraction=0.25,
         patience=patience,
         rng=42,
     )
@@ -274,7 +274,7 @@ def test_restores_best_validation_parameters(patience: int | None) -> None:
         torch.full((8, 2), 0.25),
         n_epochs=20,
         lr=0.2,
-        validation_fraction=0,
+        validation_fraction=0.25,
         patience=patience,
         rng=42,
     )
@@ -295,9 +295,12 @@ def test_restores_best_validation_parameters(patience: int | None) -> None:
         assert len(losses) == 20
 
 
+@pytest.mark.parametrize("validation_fraction", [0, 0.25])
 @pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
 def test_progress_preserves_seeded_real_zuko_training(
-    dtype: torch.dtype, capsys: pytest.CaptureFixture[str]
+    dtype: torch.dtype,
+    validation_fraction: float,
+    capsys: pytest.CaptureFixture[str],
 ) -> None:
     particles = np.random.default_rng(21).normal(size=(32, 2))
     flows = []
@@ -311,7 +314,8 @@ def test_progress_preserves_seeded_real_zuko_training(
                 batch_size=8,
                 clip_grad=0.5,
                 lr_annealing=True,
-                patience=2,
+                patience=2 if validation_fraction else None,
+                validation_fraction=validation_fraction,
                 show_progress=show_progress,
                 rng=42,
             )
@@ -321,7 +325,8 @@ def test_progress_preserves_seeded_real_zuko_training(
         output = capsys.readouterr().err
         if show_progress:
             assert "Epochs" in output
-            assert "train_loss=" in output and "val_loss=" in output
+            assert "train_loss=" in output
+            assert ("val_loss=" in output) == bool(validation_fraction)
         else:
             assert output == ""
         assert all(
@@ -336,15 +341,19 @@ def test_progress_preserves_seeded_real_zuko_training(
         )
 
 
+@pytest.mark.parametrize("validation_fraction", [0, 0.25])
 @pytest.mark.parametrize("adapter", [FlowTransform, FlowProposal])
-def test_training_controls_through_sequince(adapter: Any) -> None:
+def test_training_controls_through_sequince(
+    adapter: Any, validation_fraction: float
+) -> None:
     particles = np.random.default_rng(21).normal(size=(32, 2))
     options: dict[str, Any] = dict(
         n_epochs=3,
         batch_size=8,
         clip_grad=0.5,
         lr_annealing=True,
-        patience=2,
+        patience=2 if validation_fraction else None,
+        validation_fraction=validation_fraction,
         show_progress=False,
     )
     direct = ZukoFlow(dims=2, hidden_features=[8, 8], seed=22)
@@ -424,3 +433,51 @@ def test_rejects_invalid_training_shapes(samples: Any) -> None:
 def test_rejects_invalid_training_controls(options: dict[str, Any]) -> None:
     with pytest.raises(ValueError, match="clip_grad|patience"):
         gaussian_flow().fit(torch.zeros(8, 2), n_epochs=1, **options)
+
+
+def test_no_validation_keeps_final_parameters() -> None:
+    flow = gaussian_flow()
+    target = torch.full((2,), 0.25, dtype=torch.float64)
+    initial_rng = torch.random.get_rng_state().clone()
+    assert (
+        flow.fit(
+            target.repeat(8, 1),
+            n_epochs=20,
+            lr=0.2,
+            validation_fraction=0,
+            rng=42,
+        )
+        is flow
+    )
+    assert torch.equal(torch.random.get_rng_state(), initial_rng)
+    assert len(flow.fit_history["training_loss"]) == 20
+    assert flow.fit_history["validation_loss"] == []
+    assert flow.flow.validation_locations == []
+    assert not flow.flow.training
+
+    # Overshooting gives an earlier best epoch. With no validation, retain
+    # the final optimizer update rather than restoring that earlier location.
+    loc = torch.nn.Parameter(torch.zeros_like(target))
+    optimizer = torch.optim.Adam([loc], lr=0.2)
+    losses = []
+    for _ in range(20):
+        optimizer.zero_grad()
+        loss = (loc - target).square().sum() / 2
+        loss.backward()
+        optimizer.step()
+        losses.append(float((loc.detach() - target).square().sum()))
+    assert losses[-1] > min(losses)
+    torch.testing.assert_close(flow.flow.loc, loc, rtol=0, atol=1e-12)
+
+
+@pytest.mark.parametrize("adapter", [None, FlowTransform, FlowProposal])
+def test_patience_requires_validation(adapter: Any) -> None:
+    flow = gaussian_flow()
+    options: dict[str, Any] = {"validation_fraction": 0, "patience": 1}
+    with pytest.raises(
+        ValueError, match="patience requires validation_fraction"
+    ):
+        if adapter is None:
+            flow.fit(torch.ones(8, 2), rng=42, **options)
+        else:
+            adapter(flow, fit_kwargs=options).fit(torch.ones(8, 2), rng=42)