fix formatting

Author: priyakasimbeg

tests/test_jax_utils.py

@@ -3,18 +3,19 @@
 Run it as: pytest <path to this module>
 """
 
+from functools import partial
 import os
 
 from absl.testing import absltest
 from absl.testing import parameterized
+import flax.linen as nn
 import jax
 import jax.numpy as jnp
-import flax.linen as nn
+from jax.tree_util import tree_leaves
+from jax.tree_util import tree_map
+from jax.tree_util import tree_structure
 
-from jax.tree_util import tree_structure, tree_leaves, tree_map
 from algoperf.jax_utils import Dropout
-from functools import partial
-
 
 SEED = 1996
 DEFAULT_DROPOUT = 0.5
@@ -213,25 +214,25 @@ def test_jitted_updates(self, dropout_rate, mode):
     jitted_custom_apply = jax.jit(
         partial(cust_model.apply), static_argnames=['train'])
 
-
     def multiple_fwd_passes_custom_layer():
-        for d in [i * 0.1 * dropout_rate for i in range(0, 11)]:
-            y2 = jitted_custom_apply(
-                initial_variables_custom,
-                x,
-                train=train,
-                dropout_rate=d,
-                rngs={"dropout": dropout_rng},
-            )
-        return y2
+      for d in [i * 0.1 * dropout_rate for i in range(0, 11)]:
+        y2 = jitted_custom_apply(
+            initial_variables_custom,
+            x,
+            train=train,
+            dropout_rate=d,
+            rngs={"dropout": dropout_rng},
+        )
+      return y2
 
     def multiple_fwd_passes_original_layer():
-        for d in [i * 0.1 * dropout_rate for i in range(0, 11)]:
-             y1 = jitted_original_apply(
+      for d in [i * 0.1 * dropout_rate for i in range(0, 11)]:
+        y1 = jitted_original_apply(
             initial_variables_original,
             x,
             train=train,
             rngs={"dropout": dropout_rng})
 
+
 if __name__ == "__main__":
   absltest.main()