MAINT: update for transform API changes

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)

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
 

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)

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

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)

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