Format test_flowmatching.py for improved readability

Author: kazewong

test/unit/test_flowmatching.py

@@ -13,26 +13,31 @@
 import equinox as eqx
 import optax
 
+
 def get_simple_mlp(n_input, n_hidden, n_output, key):
     """Simple 2-layer MLP for testing."""
-    shape = [n_input] + ([n_hidden] if isinstance(n_hidden, int) else list(n_hidden)) + [n_output]
-    return MLP(
-        shape=shape,
-        key=key,
-        activation=jax.nn.swish
+    shape = (
+        [n_input]
+        + ([n_hidden] if isinstance(n_hidden, int) else list(n_hidden))
+        + [n_output]
     )
+    return MLP(shape=shape, key=key, activation=jax.nn.swish)
+
 
 ##############################
 # Solver Tests
 ##############################
 
+
 class TestSolver:
     @pytest.fixture
     def solver(self):
         key = jax.random.PRNGKey(0)
         n_dim = 3
         n_hidden = 4
-        mlp = get_simple_mlp(n_input=n_dim+1, n_hidden=n_hidden, n_output=n_dim, key=key)
+        mlp = get_simple_mlp(
+            n_input=n_dim + 1, n_hidden=n_hidden, n_output=n_dim, key=key
+        )
         return Solver(model=mlp, method=Dopri5()), key, n_dim
 
     def test_sample_shape_and_finiteness(self, solver):
@@ -57,10 +62,12 @@ def test_sample_various_dt(self, solver, dt):
         assert samples.shape == (3, n_dim)
         assert jnp.isfinite(samples).all()
 
+
 ##############################
 # Path & Scheduler Tests
 ##############################
 
+
 class TestPathAndScheduler:
     def test_path_sample_shapes_and_values(self):
         n_dim = 2
@@ -82,25 +89,29 @@ def test_condotscheduler_call_output(self, t):
         assert len(out) == 4
         assert all(isinstance(float(x), float) for x in out)
 
+
 ##############################
 # FlowMatchingModel Tests
 ##############################
 
+
 class TestFlowMatchingModel:
     @pytest.fixture
     def model(self):
         key = jax.random.PRNGKey(42)
         n_dim = 2
         n_hidden = 8
-        mlp = get_simple_mlp(n_input=n_dim+1, n_hidden=n_hidden, n_output=n_dim, key=key)
+        mlp = get_simple_mlp(
+            n_input=n_dim + 1, n_hidden=n_hidden, n_output=n_dim, key=key
+        )
         solver = Solver(model=mlp, method=Dopri5())
         scheduler = CondOTScheduler()
         path = Path(scheduler=scheduler)
         model = FlowMatchingModel(
             solver=solver,
             path=path,
             data_mean=jnp.zeros(n_dim),
-            data_cov=jnp.eye(n_dim)
+            data_cov=jnp.eye(n_dim),
         )
         return model, key, n_dim
 
@@ -130,21 +141,27 @@ def test_log_prob_edge_cases(self, model):
             logp = model.log_prob(arr)
             logp_arr = jnp.asarray(logp)
             assert logp_arr.size == 1
-            assert jnp.isfinite(logp_arr).all() or jnp.isnan(logp_arr).all()  # may be nan for extreme values
+            assert (
+                jnp.isfinite(logp_arr).all() or jnp.isnan(logp_arr).all()
+            )  # may be nan for extreme values
 
     def test_save_and_load(self, tmp_path, model):
         model, key, n_dim = model
         save_path = str(tmp_path / "test_model")
         model.save_model(save_path)
         loaded = model.load_model(save_path)
         x = jax.random.normal(key, (2, n_dim))
-        assert jnp.allclose(eqx.filter_vmap(model.log_prob)(x), eqx.filter_vmap(loaded.log_prob)(x))
+        assert jnp.allclose(
+            eqx.filter_vmap(model.log_prob)(x), eqx.filter_vmap(loaded.log_prob)(x)
+        )
 
     def test_properties(self, model):
         model, key, n_dim = model
         mean = jnp.arange(n_dim)
         cov = jnp.eye(n_dim) * 2
-        model2 = FlowMatchingModel(solver=model.solver, path=model.path, data_mean=mean, data_cov=cov)
+        model2 = FlowMatchingModel(
+            solver=model.solver, path=model.path, data_mean=mean, data_cov=cov
+        )
         assert model2.n_features == n_dim
         assert jnp.allclose(model2.data_mean, mean)
         assert jnp.allclose(model2.data_cov, cov)
@@ -157,7 +174,6 @@ def test_print_parameters_notimplemented(self, model):
     def test_train_step_and_epoch(self, model):
         model, key, n_dim = model
         n_batch = 5
-        n_hidden = 8
         x0 = jax.random.normal(key, (n_batch, n_dim))
         x1 = jax.random.normal(key, (n_batch, n_dim))
         t = jax.random.uniform(key, (n_batch, 1))
@@ -170,6 +186,8 @@ def test_train_step_and_epoch(self, model):
         assert jnp.isfinite(loss)
         assert isinstance(model2, FlowMatchingModel)
         data = (x0, x1, t)
-        loss_epoch, model3, state3 = model.train_epoch(key, optim, state, data, batch_size=n_batch)
+        loss_epoch, model3, state3 = model.train_epoch(
+            key, optim, state, data, batch_size=n_batch
+        )
         assert jnp.isfinite(loss_epoch)
         assert isinstance(model3, FlowMatchingModel)