From 8b1e2aa7b04bc55e1e3884e3e7556b654fb6977a Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Fri, 21 Aug 2026 11:56:28 +0200 Subject: [PATCH 1/2] MAINT: update for transform API changes --- src/sequince_density/flows/base.py | 20 +++++++++++++------- src/sequince_density/flows/flowjax.py | 8 +++++--- src/sequince_density/flows/zuko.py | 14 +++++++++----- 3 files changed, 27 insertions(+), 15 deletions(-) diff --git a/src/sequince_density/flows/base.py b/src/sequince_density/flows/base.py index 7e75582..31a9fc8 100644 --- a/src/sequince_density/flows/base.py +++ b/src/sequince_density/flows/base.py @@ -53,18 +53,24 @@ def sample(self, n_samples: int, context=None, *, rng=None): def sample_and_log_prob(self, n_samples: int, context=None, *, rng=None): raise NotImplementedError - def fit_particle_transform(self, x, context=None): + def fit_particle_transform(self, x, context=None, *, rng=None): self.particle_transform = self.particle_transform.fit( - x, context=context + x, context=context, rng=rng + ) + transformed, _, _ = self.particle_transform.forward( + x, context=context, rng=rng ) - transformed, _ = self.particle_transform.forward(x, context=context) return transformed - def forward_particle_transform(self, x, context=None): - return self.particle_transform.forward(x, context=context) + def forward_particle_transform(self, x, context=None, *, rng=None): + return self.particle_transform.forward(x, context=context, rng=rng) - def inverse_particle_transform(self, x, context=None): - return self.particle_transform.inverse(x, context=context) + def inverse_particle_transform( + self, x, context=None, *, auxiliary=None, rng=None + ): + return self.particle_transform.inverse( + x, context=context, auxiliary=auxiliary, rng=rng + ) def __new__(cls, *args, **kwargs): obj = super().__new__(cls) diff --git a/src/sequince_density/flows/flowjax.py b/src/sequince_density/flows/flowjax.py index 120e0f4..a997118 100644 --- a/src/sequince_density/flows/flowjax.py +++ b/src/sequince_density/flows/flowjax.py @@ -74,7 +74,7 @@ def flow(self): def fit(self, samples, context=None, *, rng=None, **kwargs) -> FlowJax: values = jnp.asarray(samples, dtype=self.dtype) transformed = jnp.asarray( - self.fit_particle_transform(values, context=context), + self.fit_particle_transform(values, context=context, rng=rng), dtype=self.dtype, ) fit_key = self._random_key(rng) @@ -115,7 +115,7 @@ def load_state_dict(self, state): def forward(self, samples, context=None): values = jnp.asarray(samples, dtype=self.dtype) - transformed, transform_log_det = self.forward_particle_transform( + transformed, transform_log_det, _ = self.forward_particle_transform( values, context=context, ) @@ -137,7 +137,7 @@ def inverse(self, samples, context=None): def log_prob(self, samples, context=None): values = jnp.asarray(samples, dtype=self.dtype) - transformed, log_det = self.forward_particle_transform( + transformed, log_det, _ = self.forward_particle_transform( values, context=context, ) @@ -158,6 +158,7 @@ def sample(self, n_samples: int, context=None, *, rng=None): values, _ = self.inverse_particle_transform( transformed, context=context, + rng=rng, ) return values @@ -170,6 +171,7 @@ def sample_and_log_prob(self, n_samples: int, context=None, *, rng=None): values, log_det = self.inverse_particle_transform( transformed, context=context, + rng=rng, ) return values, log_prob - log_det diff --git a/src/sequince_density/flows/zuko.py b/src/sequince_density/flows/zuko.py index 458704d..8676220 100644 --- a/src/sequince_density/flows/zuko.py +++ b/src/sequince_density/flows/zuko.py @@ -84,6 +84,7 @@ def fit( lr=lr, batch_size=batch_size, validation_fraction=validation_fraction, + rng=rng, ) def _fit( @@ -94,9 +95,10 @@ def _fit( lr: float = 1e-3, batch_size: int = 500, validation_fraction: float = 0.2, + rng=None, ) -> "ZukoFlow": x = self._as_torch(samples) - x_prime = self.fit_particle_transform(x, context=context) + x_prime = self.fit_particle_transform(x, context=context, rng=rng) indices = torch.randperm(x_prime.shape[0], device=x_prime.device) x_prime = x_prime[indices] @@ -186,26 +188,28 @@ def sample_and_log_prob(self, n_samples: int, context=None, *, rng=None): with torch_random_context(rng, device=self.device), torch.no_grad(): x_prime, log_prob = self.flow().rsample_and_log_prob((n_samples,)) x, log_abs_det_jacobian = self.inverse_particle_transform( - x_prime, context=context + x_prime, context=context, rng=rng ) return x, log_prob - log_abs_det_jacobian def sample(self, n_samples: int, context=None, *, rng=None): with torch_random_context(rng, device=self.device), torch.no_grad(): x_prime = self.flow().rsample((n_samples,)) - x, _ = self.inverse_particle_transform(x_prime, context=context) + x, _ = self.inverse_particle_transform( + x_prime, context=context, rng=rng + ) return x def log_prob(self, x, context=None): x = self._as_torch(x) - x_prime, log_abs_det_jacobian = self.forward_particle_transform( + x_prime, log_abs_det_jacobian, _ = self.forward_particle_transform( x, context=context ) return self.flow().log_prob(x_prime) + log_abs_det_jacobian def forward(self, x, context=None): x = self._as_torch(x) - x_prime, log_j_rescale = self.forward_particle_transform( + x_prime, log_j_rescale, _ = self.forward_particle_transform( x, context=context ) z, log_abs_det_jacobian = self.flow().transform.call_and_ladj(x_prime) From 5f9c75c20dbeecac6ad09521ed537a74cf19b078 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Fri, 21 Aug 2026 12:02:12 +0200 Subject: [PATCH 2/2] TST: add tests for auxiliary and rng --- tests/test_flow_base.py | 58 +++++++++++++++++++++++++++++++++++++++++ tests/test_flowjax.py | 6 ++++- tests/test_zuko.py | 12 ++++++--- 3 files changed, 72 insertions(+), 4 deletions(-) create mode 100644 tests/test_flow_base.py diff --git a/tests/test_flow_base.py b/tests/test_flow_base.py new file mode 100644 index 0000000..8f5bb80 --- /dev/null +++ b/tests/test_flow_base.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +import numpy as np +from orng import RandomGenerator +from sequince.transforms import ParticleTransform + +from sequince_density.flows.base import Flow + + +class RecordingTransform(ParticleTransform): + def __init__(self) -> None: + self.calls: list[tuple] = [] + + def fit(self, particles, context=None, *, rng=None): + self.calls.append(("fit", context, rng, None)) + return self + + def forward(self, particles, context=None, *, rng=None): + auxiliary = {"mode": np.ones_like(particles)} + self.calls.append(("forward", context, rng, auxiliary)) + return particles + 1.0, np.zeros(particles.shape[0]), auxiliary + + def inverse(self, values, context=None, *, auxiliary=None, rng=None): + self.calls.append(("inverse", context, rng, auxiliary)) + return values - 1.0, np.zeros(values.shape[0]) + + +def test_particle_transform_helpers_preserve_auxiliary_and_rng() -> None: + particles = np.zeros((3, 2)) + context = object() + rng = RandomGenerator(seed=12) + transform = RecordingTransform() + flow = Flow(dims=2, particle_transform=transform) + + fitted = flow.fit_particle_transform(particles, context=context, rng=rng) + transformed, log_det, auxiliary = flow.forward_particle_transform( + particles, context=context, rng=rng + ) + restored, inverse_log_det = flow.inverse_particle_transform( + transformed, + context=context, + auxiliary=auxiliary, + rng=rng, + ) + + np.testing.assert_array_equal(fitted, particles + 1.0) + np.testing.assert_array_equal(restored, particles) + np.testing.assert_array_equal(log_det, 0.0) + np.testing.assert_array_equal(inverse_log_det, 0.0) + assert [call[0] for call in transform.calls] == [ + "fit", + "forward", + "forward", + "inverse", + ] + assert all(call[1] is context for call in transform.calls) + assert all(call[2] is rng for call in transform.calls) + assert transform.calls[-1][3] is auxiliary diff --git a/tests/test_flowjax.py b/tests/test_flowjax.py index f6185d8..5ab544a 100644 --- a/tests/test_flowjax.py +++ b/tests/test_flowjax.py @@ -2,6 +2,7 @@ import numpy as np import pytest +from orng import RandomGenerator jax = pytest.importorskip("jax") jnp = pytest.importorskip("jax.numpy") @@ -26,13 +27,16 @@ def test_flowjax_fits_samples_and_round_trips() -> None: returned = flow.fit( particles, + rng=RandomGenerator(seed=19), max_epochs=1, batch_size=16, show_progress=False, ) latent, forward_log_det = flow.forward(particles[:4]) restored, inverse_log_det = flow.inverse(latent) - samples, log_prob = flow.sample_and_log_prob(5) + samples, log_prob = flow.sample_and_log_prob( + 5, rng=RandomGenerator(seed=20) + ) assert returned is flow assert isinstance(samples, jax.Array) diff --git a/tests/test_zuko.py b/tests/test_zuko.py index e5162cf..b8fbc42 100644 --- a/tests/test_zuko.py +++ b/tests/test_zuko.py @@ -2,6 +2,7 @@ import numpy as np import pytest +from orng import RandomGenerator from sequince.transforms import ( AffineTransform, CompositeTransform, @@ -36,9 +37,14 @@ def test_zuko_flow_transform_fits_and_round_trips() -> None: ), ] ) - transform = transform.fit(particles) - latent, log_det = transform.forward(particles) - restored, inverse_log_det = transform.inverse(latent) + transform_rng = RandomGenerator(seed=19) + transform = transform.fit(particles, rng=transform_rng) + latent, log_det, auxiliary = transform.forward( + particles, rng=transform_rng + ) + restored, inverse_log_det = transform.inverse( + latent, auxiliary=auxiliary, rng=transform_rng + ) assert isinstance(latent, np.ndarray) assert latent.shape == particles.shape