lint

aleatory-science · Dec 4, 2023 · ce45e56 · ce45e56
1 parent 783f9f3
commit ce45e56
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/numpyro/contrib/control_flow/scan.py b/numpyro/contrib/control_flow/scan.py
@@ -18,8 +18,11 @@
 def _replay_wrapper(replay_trace, trace, i, length):
     def get_ith_value(site):
         site_len = jnp.shape(site["value"])[0]
-
-        if site["name"] not in trace or site_len != length or site["type"] not in ("sample", "deterministic"):
+        if (
+            site["name"] not in trace
+            or site_len != length
+            or site["type"] not in ("sample", "deterministic")
+        ):
             return site
 
         site["value"] = site["value"][i]