Fix yapf

Author: runame

algorithmic_efficiency/workloads/imagenet_vit/imagenet_jax/workload.py

@@ -21,8 +21,7 @@
 # Make sure we inherit from the ViT base workload first.
 class ImagenetVitWorkload(BaseImagenetVitWorkload, ImagenetResNetWorkload):
 
-  def initialized(self,
-                  key: spec.RandomState,
+  def initialized(self, key: spec.RandomState,
                   model: nn.Module) -> spec.ModelInitState:
     input_shape = (1, 224, 224, 3)
     params_rng, dropout_rng = jax.random.split(key)

algorithmic_efficiency/workloads/wmt/wmt_jax/workload.py

@@ -235,10 +235,10 @@ def init_model_fn(
     eval_config = replace(model_config, deterministic=True)
     self._eval_model = models.Transformer(eval_config)
     params_rng, dropout_rng = jax.random.split(rng)
-    initial_variables = jax.jit(self._eval_model.init)(
-        {'params': params_rng, 'dropout': dropout_rng},
-        jnp.ones(input_shape, jnp.float32),
-        jnp.ones(target_shape, jnp.float32))
+    initial_variables = jax.jit(
+        self._eval_model.init)({'params': params_rng, 'dropout': dropout_rng},
+                               jnp.ones(input_shape, jnp.float32),
+                               jnp.ones(target_shape, jnp.float32))
 
     initial_params = initial_variables['params']
     self._param_shapes = param_utils.jax_param_shapes(initial_params)