ci: apply fixes

Signed-off-by: nstarman <[email protected]>

Author: nstarman

numpyro/contrib/control_flow/scan.py

@@ -197,9 +197,10 @@ def body_fn(wrapped_carry, x, prefix=None):
             )
         return (i + 1, rng_key, new_carry), (PytreeTrace(trace), y)
 
-    with handlers.block(
-        hide_fn=lambda site: not site["name"].startswith("_PREV_")
-    ), enum(first_available_dim=first_available_dim):
+    with (
+        handlers.block(hide_fn=lambda site: not site["name"].startswith("_PREV_")),
+        enum(first_available_dim=first_available_dim),
+    ):
         wrapped_carry = (0, rng_key, init)
         y0s = []
         # We run unroll_steps + 1 where the last step is used for rolling with `lax.scan`

numpyro/infer/autoguide.py

@@ -1463,8 +1463,10 @@ def _sample_latent(self, *args, **kwargs):
         if self.global_guide is not None:
             global_latents = self.global_guide(*args, **kwargs)
             rng_key = numpyro.prng_key()
-            with handlers.block(), handlers.seed(rng_seed=rng_key), handlers.substitute(
-                data=global_latents
+            with (
+                handlers.block(),
+                handlers.seed(rng_seed=rng_key),
+                handlers.substitute(data=global_latents),
             ):
                 global_outputs = self.global_guide.model(*args, **kwargs)
             local_args = (global_outputs,)
@@ -1575,9 +1577,12 @@ def fn(x):
             if self.local_guide is not None:
                 key = numpyro.prng_key()
                 subsample_guide = partial(_subsample_model, self.local_guide)
-                with handlers.block(), handlers.trace() as tr, handlers.seed(
-                    rng_seed=key
-                ), handlers.substitute(data=local_guide_params):
+                with (
+                    handlers.block(),
+                    handlers.trace() as tr,
+                    handlers.seed(rng_seed=key),
+                    handlers.substitute(data=local_guide_params),
+                ):
                     with warnings.catch_warnings():
                         warnings.simplefilter("ignore")
                         subsample_guide(*local_args, **local_kwargs)

numpyro/infer/elbo.py

@@ -874,8 +874,11 @@ def get_importance_trace_enum(
         trace as _trace,
     )
 
-    with plate_to_enum_plate(), enum(
-        first_available_dim=(-max_plate_nesting - 1) if max_plate_nesting else None
+    with (
+        plate_to_enum_plate(),
+        enum(
+            first_available_dim=(-max_plate_nesting - 1) if max_plate_nesting else None
+        ),
     ):
         guide = substitute(guide, data=params)
         with _without_rsample_stop_gradient():

numpyro/infer/inspect.py

@@ -58,8 +58,10 @@ def get_trace():
 
 def _get_log_probs(model, model_args, model_kwargs, **sample):
     # Note: We use seed 0 for parameter initialization.
-    with handlers.trace() as tr, handlers.seed(rng_seed=0), handlers.substitute(
-        data=sample
+    with (
+        handlers.trace() as tr,
+        handlers.seed(rng_seed=0),
+        handlers.substitute(data=sample),
     ):
         model(*model_args, **model_kwargs)
     return {
@@ -370,8 +372,9 @@ def process_message(self, msg):
 
         # Note: We use seed 0 for parameter initialization.
         with handlers.trace() as tr, handlers.seed(rng_seed=0):
-            with handlers.substitute(data=sample), substitute_deterministic(
-                data=sample
+            with (
+                handlers.substitute(data=sample),
+                substitute_deterministic(data=sample),
             ):
                 model(*model_args, **model_kwargs)
         provenance_arrays = {}

test/infer/test_hmc_util.py

@@ -57,8 +57,9 @@ def optimize(f):
 @pytest.mark.parametrize("regularize", [True, False])
 @pytest.mark.filterwarnings("ignore:numpy.linalg support is experimental:UserWarning")
 def test_welford_covariance(jitted, diagonal, regularize):
-    with optional(jitted, disable_jit()), optional(
-        jitted, control_flow_prims_disabled()
+    with (
+        optional(jitted, disable_jit()),
+        optional(jitted, control_flow_prims_disabled()),
     ):
         np.random.seed(0)
         loc = np.random.randn(3)

test/test_distributions.py

@@ -2221,8 +2221,12 @@ def g(x):
 
 
 def test_beta_proportion_invalid_mean():
-    with dist.distribution.validation_enabled(), pytest.raises(
-        ValueError, match=r"^BetaProportion distribution got invalid mean parameter\.$"
+    with (
+        dist.distribution.validation_enabled(),
+        pytest.raises(
+            ValueError,
+            match=r"^BetaProportion distribution got invalid mean parameter\.$",
+        ),
     ):
         dist.BetaProportion(1.0, 1.0)
 

test/test_handlers.py

@@ -372,8 +372,10 @@ def test_subsample_substitute():
     data = jnp.arange(100.0)
     subsample_size = 7
     subsample = jnp.array([13, 3, 30, 4, 1, 68, 5])
-    with handlers.trace() as tr, handlers.seed(rng_seed=0), handlers.substitute(
-        data={"a": subsample}
+    with (
+        handlers.trace() as tr,
+        handlers.seed(rng_seed=0),
+        handlers.substitute(data={"a": subsample}),
     ):
         with numpyro.plate("a", len(data), subsample_size=subsample_size) as idx:
             assert data[idx].shape == (subsample_size,)