diff --git a/CHANGELOG.md b/CHANGELOG.md
index af05bb3..e29ba76 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,11 +8,21 @@ The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
### Added
+- 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
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 94ba6ad..fc23565 100644
--- a/docs/normalizing_flows.rst
+++ b/docs/normalizing_flows.rst
@@ -34,6 +34,41 @@ Estimation*. Credit the architecture you select as well as the backend
implementation: `Zuko `_ or
`FlowJAX `_.
+Torch training controls
+-----------------------
+
+``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
+
+ flow.fit(
+ particles, n_epochs=100, clip_grad=5.0, lr_annealing=True,
+ patience=10, show_progress=True, rng=42,
+ )
+
+Pass the same options through ``fit_kwargs`` in ``FlowTransform`` or
+``FlowProposal``. See :py:meth:`~sequince_density.flows.zuko.ZukoFlow.fit`
+for all training defaults.
+
+``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
+identify the stage and affected dimensions (zero-based).
+
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..9840f79 100644
--- a/src/sequince_density/flows/zuko.py
+++ b/src/sequince_density/flows/zuko.py
@@ -8,8 +8,9 @@
import array_api_compat.torch as torch_api
import torch
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__)
@@ -73,9 +74,54 @@ 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, restoring the best epoch when validating.
+
+ 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. 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.
+ 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. 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
+ RNG for fitting, isolated from global Torch state when supplied.
+
+ Returns
+ -------
+ ZukoFlow
+ This fitted adapter. fit_history contains losses for every
+ completed epoch; validation_loss is empty when validation is off.
+ """
with torch_random_context(rng, device=self.device):
return self._fit(
samples,
@@ -84,6 +130,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,15 +145,36 @@ 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.")
+ 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)
+ 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)
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:
@@ -115,47 +186,83 @@ 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)
+ 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()
+ 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
+ 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)
+ 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)
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)
diff --git a/tests/test_zuko.py b/tests/test_zuko.py
index 0f3662c..e6b9b0b 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,337 @@ 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.25,
+ 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.25,
+ 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.25,
+ 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.25,
+ 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.25,
+ 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("validation_fraction", [0, 0.25])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+def test_progress_preserves_seeded_real_zuko_training(
+ dtype: torch.dtype,
+ validation_fraction: float,
+ 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 if validation_fraction else None,
+ validation_fraction=validation_fraction,
+ 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
+ assert ("val_loss=" in output) == bool(validation_fraction)
+ 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("validation_fraction", [0, 0.25])
+@pytest.mark.parametrize("adapter", [FlowTransform, FlowProposal])
+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 if validation_fraction else None,
+ validation_fraction=validation_fraction,
+ 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)
+
+
+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)