fix formatting for ruff

Author: priyakasimbeg

algoperf/jax_sharding_utils.py

@@ -20,7 +20,8 @@ def shard_along_batch_dim(x):
   """Shards a tensor across all devices."""
   mesh = jax.sharding.Mesh(jax.devices(), ('batch',))
   return jax.tree.map(
-      lambda x: jax.device_put(x, NamedSharding(mesh, P('batch'))), x)
+    lambda x: jax.device_put(x, NamedSharding(mesh, P('batch'))), x
+  )
 
 
 def replicate(x):
@@ -32,5 +33,7 @@ def replicate(x):
 def display_shard_info(x: jax.Array):
   """Displays shard info of a jax array."""
   for shard in x.addressable_shards:
-    print(f"shard.device: {shard.device}, index: {shard.index}, replica_id:"
-          f" {shard.replica_id}.\n")
+    print(
+      f'shard.device: {shard.device}, index: {shard.index}, replica_id:'
+      f' {shard.replica_id}.\n'
+    )

algoperf/workloads/cifar/cifar_jax/workload.py

@@ -176,35 +176,38 @@ def _compute_metrics(
     return metrics
 
   def _eval_model(
-      self,
-      params: spec.ParameterContainer,
-      batch: Dict[str, spec.Tensor],
-      model_state: spec.ModelAuxiliaryState,
-      rng: spec.RandomState) -> Dict[spec.Tensor, spec.ModelAuxiliaryState]:
+    self,
+    params: spec.ParameterContainer,
+    batch: Dict[str, spec.Tensor],
+    model_state: spec.ModelAuxiliaryState,
+    rng: spec.RandomState,
+  ) -> Dict[spec.Tensor, spec.ModelAuxiliaryState]:
     """Return the mean accuracy and loss as a dict."""
 
     @functools.partial(
-        jax.jit,
-        in_shardings=(
-            jax_sharding_utils.get_replicate_sharding(),  # params
-            jax_sharding_utils.get_batch_dim_sharding(),  # batch
-            jax_sharding_utils.get_replicate_sharding(),  # model_state
-            jax_sharding_utils.get_batch_dim_sharding(),  # rng
-        ),
+      jax.jit,
+      in_shardings=(
+        jax_sharding_utils.get_replicate_sharding(),  # params
+        jax_sharding_utils.get_batch_dim_sharding(),  # batch
+        jax_sharding_utils.get_replicate_sharding(),  # model_state
+        jax_sharding_utils.get_batch_dim_sharding(),  # rng
+      ),
     )
     def _eval_model_jitted(
-        params: spec.ParameterContainer,
-        batch: Dict[str, spec.Tensor],
-        model_state: spec.ModelAuxiliaryState,
-        rng: spec.RandomState) -> Dict[spec.Tensor, spec.ModelAuxiliaryState]:
+      params: spec.ParameterContainer,
+      batch: Dict[str, spec.Tensor],
+      model_state: spec.ModelAuxiliaryState,
+      rng: spec.RandomState,
+    ) -> Dict[spec.Tensor, spec.ModelAuxiliaryState]:
       """Return the mean accuracy and loss as a dict."""
       logits, _ = self.model_fn(
-          params,
-          batch,
-          model_state,
-          spec.ForwardPassMode.EVAL,
-          rng,
-          update_batch_norm=False)
+        params,
+        batch,
+        model_state,
+        spec.ForwardPassMode.EVAL,
+        rng,
+        update_batch_norm=False,
+      )
       weights = batch.get('weights')
       if weights is None:
         weights = jnp.ones(len(logits))

algoperf/workloads/criteo1tb/criteo1tb_jax/workload.py

@@ -134,16 +134,17 @@ def model_fn(
     return logits_batch, None
 
   @functools.partial(
-      jax.jit,
-      in_shardings=(
-          jax_sharding_utils.get_replicate_sharding(),
-          jax_sharding_utils.get_batch_dim_sharding(),
-      ),
-      static_argnums=(0,),
-      out_shardings=jax_sharding_utils.get_replicate_sharding())
-  def _eval_batch_jitted(self,
-                         params: spec.ParameterContainer,
-                         batch: Dict[str, spec.Tensor]) -> spec.Tensor:
+    jax.jit,
+    in_shardings=(
+      jax_sharding_utils.get_replicate_sharding(),
+      jax_sharding_utils.get_batch_dim_sharding(),
+    ),
+    static_argnums=(0,),
+    out_shardings=jax_sharding_utils.get_replicate_sharding(),
+  )
+  def _eval_batch_jitted(
+    self, params: spec.ParameterContainer, batch: Dict[str, spec.Tensor]
+  ) -> spec.Tensor:
     logits, _ = self.model_fn(
       params,
       batch,

algoperf/workloads/fastmri/fastmri_jax/workload.py

@@ -101,16 +101,21 @@ def loss_fn(
     }
 
   @functools.partial(
-      jax.jit,
-      in_shardings=(jax_sharding_utils.get_replicate_sharding(),
-                    jax_sharding_utils.get_batch_dim_sharding(),
-                    jax_sharding_utils.get_replicate_sharding()),
-      static_argnums=(0,),
-      out_shardings=jax_sharding_utils.get_replicate_sharding())
-  def _eval_model(self,
-                  params: spec.Tensor,
-                  batch: Dict[str, spec.Tensor],
-                  rng: spec.RandomState) -> Dict[str, spec.Tensor]:
+    jax.jit,
+    in_shardings=(
+      jax_sharding_utils.get_replicate_sharding(),
+      jax_sharding_utils.get_batch_dim_sharding(),
+      jax_sharding_utils.get_replicate_sharding(),
+    ),
+    static_argnums=(0,),
+    out_shardings=jax_sharding_utils.get_replicate_sharding(),
+  )
+  def _eval_model(
+    self,
+    params: spec.Tensor,
+    batch: Dict[str, spec.Tensor],
+    rng: spec.RandomState,
+  ) -> Dict[str, spec.Tensor]:
     """Return the SSIM and loss as a dict."""
     logits, _ = self.model_fn(
       params,
@@ -166,13 +171,13 @@ def _eval_model_on_split(
         num_batches=num_batches,
       )
 
-    total_metrics = {'ssim': 0., 'loss': 0.}
+    total_metrics = {'ssim': 0.0, 'loss': 0.0}
     for _ in range(num_batches):
       batch = next(self._eval_iters[split])
       # We already sum these metrics across devices inside _eval_model.
       synced_metrics = self._eval_model(params, batch, model_rng)
       total_metrics = {
-          k: v + synced_metrics[k] for k, v in total_metrics.items()
+        k: v + synced_metrics[k] for k, v in total_metrics.items()
       }
     return {k: float(v.item() / num_examples) for k, v in total_metrics.items()}
 

algoperf/workloads/imagenet_resnet/imagenet_jax/workload.py

@@ -107,33 +107,36 @@ def init_model_fn(
     self._param_shapes = param_utils.jax_param_shapes(params)
     self._param_types = param_utils.jax_param_types(self._param_shapes)
     params = jax.tree.map(
-        lambda x: jax.device_put(x, jax_sharding_utils.get_replicate_sharding()
-                                ),
-        params)
+      lambda x: jax.device_put(x, jax_sharding_utils.get_replicate_sharding()),
+      params,
+    )
     model_state = jax.tree.map(
-        lambda x: jax.device_put(x, jax_sharding_utils.get_replicate_sharding()
-                                ),
-        model_state)
+      lambda x: jax.device_put(x, jax_sharding_utils.get_replicate_sharding()),
+      model_state,
+    )
     return params, model_state
 
   def is_output_params(self, param_key: spec.ParameterKey) -> bool:
     return param_key == 'Dense_0'
 
   @functools.partial(
-      jax.jit,
-      in_shardings=(
-          jax_sharding_utils.get_replicate_sharding(),  # params
-          jax_sharding_utils.get_batch_dim_sharding(),  # batch
-          jax_sharding_utils.get_replicate_sharding(),  # model_state
-          jax_sharding_utils.get_replicate_sharding(),  # rng
-      ),
-      static_argnums=(0,),
-      out_shardings=jax_sharding_utils.get_replicate_sharding())
-  def _eval_model(self,
-                  params: spec.ParameterContainer,
-                  batch: Dict[str, spec.Tensor],
-                  model_state: spec.ModelAuxiliaryState,
-                  rng: spec.RandomState) -> Dict[str, spec.Tensor]:
+    jax.jit,
+    in_shardings=(
+      jax_sharding_utils.get_replicate_sharding(),  # params
+      jax_sharding_utils.get_batch_dim_sharding(),  # batch
+      jax_sharding_utils.get_replicate_sharding(),  # model_state
+      jax_sharding_utils.get_replicate_sharding(),  # rng
+    ),
+    static_argnums=(0,),
+    out_shardings=jax_sharding_utils.get_replicate_sharding(),
+  )
+  def _eval_model(
+    self,
+    params: spec.ParameterContainer,
+    batch: Dict[str, spec.Tensor],
+    model_state: spec.ModelAuxiliaryState,
+    rng: spec.RandomState,
+  ) -> Dict[str, spec.Tensor]:
     logits, _ = self.model_fn(
       params,
       batch,

algoperf/workloads/imagenet_vit/imagenet_jax/workload.py

@@ -10,7 +10,7 @@
 
 from algoperf import param_utils, spec, jax_sharding_utils
 from algoperf.workloads.imagenet_resnet.imagenet_jax.workload import (
-    ImagenetResNetWorkload,
+  ImagenetResNetWorkload,
 )
 from algoperf.workloads.imagenet_vit.imagenet_jax import models
 from algoperf.workloads.imagenet_vit.workload import (

algoperf/workloads/librispeech_conformer/librispeech_jax/workload.py

@@ -15,8 +15,9 @@
 from algoperf import spec
 from algoperf.workloads.librispeech_conformer import metrics
 from algoperf.workloads.librispeech_conformer import workload
-from algoperf.workloads.librispeech_conformer.input_pipeline import \
-    LibriSpeechDataset
+from algoperf.workloads.librispeech_conformer.input_pipeline import (
+  LibriSpeechDataset,
+)
 from algoperf.workloads.librispeech_conformer.librispeech_jax import models
 
 
@@ -317,21 +318,23 @@ def greedy_decode(
     return hyp, hyp_paddings
 
   @functools.partial(
-      jax.jit,
-      in_shardings=(
-          jax_sharding_utils.get_replicate_sharding(),  # params
-          jax_sharding_utils.get_batch_dim_sharding(),  # batch
-          jax_sharding_utils.get_replicate_sharding(),  # model_state
-          jax_sharding_utils.get_replicate_sharding(),  # rng
-      ),
-      out_shardings=jax_sharding_utils.get_batch_dim_sharding(),
-      static_argnums=(0,))
+    jax.jit,
+    in_shardings=(
+      jax_sharding_utils.get_replicate_sharding(),  # params
+      jax_sharding_utils.get_batch_dim_sharding(),  # batch
+      jax_sharding_utils.get_replicate_sharding(),  # model_state
+      jax_sharding_utils.get_replicate_sharding(),  # rng
+    ),
+    out_shardings=jax_sharding_utils.get_batch_dim_sharding(),
+    static_argnums=(0,),
+  )
   def _eval_step(
-      self,
-      params: spec.ParameterContainer,
-      batch: Dict[str, spec.Tensor],
-      model_state: spec.ModelAuxiliaryState,
-      rng: spec.RandomState) -> Tuple[spec.Tensor, spec.ModelAuxiliaryState]:
+    self,
+    params: spec.ParameterContainer,
+    batch: Dict[str, spec.Tensor],
+    model_state: spec.ModelAuxiliaryState,
+    rng: spec.RandomState,
+  ) -> Tuple[spec.Tensor, spec.ModelAuxiliaryState]:
     (logits, logit_paddings), _ = self.model_fn(
       params,
       batch,
@@ -346,40 +349,38 @@ def _eval_step(
     targets, target_paddings = batch['targets']
     # Convert metrics bundle to dictionary
     metrics_dict = {
-        'loss_per_example':
-            loss['per_example'],
-        'decoded':
-            decoded,
-        'decoded_paddings':
-            decoded_paddings,
-        'targets':
-            targets,
-        'target_paddings':
-            target_paddings,
-        'n_valid_examples':
-            jnp.zeros((len(jax.devices()), 1)) + loss['n_valid_examples']
+      'loss_per_example': loss['per_example'],
+      'decoded': decoded,
+      'decoded_paddings': decoded_paddings,
+      'targets': targets,
+      'target_paddings': target_paddings,
+      'n_valid_examples': jnp.zeros((len(jax.devices()), 1))
+      + loss['n_valid_examples'],
     }
     return metrics_dict
 
-  def eval_step(self,
-                params: spec.ParameterContainer,
-                batch: Dict[str, spec.Tensor],
-                model_state: spec.ModelAuxiliaryState,
-                rng: spec.RandomState):
+  def eval_step(
+    self,
+    params: spec.ParameterContainer,
+    batch: Dict[str, spec.Tensor],
+    model_state: spec.ModelAuxiliaryState,
+    rng: spec.RandomState,
+  ):
     """Evaluates the model and returns a metrics bundle."""
     metrics_dict = self._eval_step(params, batch, model_state, rng)
 
     # Convert dictionary back to metrics bundle
     metrics_bundle = self.metrics_bundle.single_from_model_output(
-        loss_dict={
-            'summed': metrics_dict['loss_per_example'].sum(),
-            'per_example': metrics_dict['loss_per_example'],
-            'n_valid_examples': metrics_dict['n_valid_examples'].sum()
-        },
-        decoded=metrics_dict['decoded'],
-        decoded_paddings=metrics_dict['decoded_paddings'],
-        targets=metrics_dict['targets'],
-        target_paddings=metrics_dict['target_paddings'])
+      loss_dict={
+        'summed': metrics_dict['loss_per_example'].sum(),
+        'per_example': metrics_dict['loss_per_example'],
+        'n_valid_examples': metrics_dict['n_valid_examples'].sum(),
+      },
+      decoded=metrics_dict['decoded'],
+      decoded_paddings=metrics_dict['decoded_paddings'],
+      targets=metrics_dict['targets'],
+      target_paddings=metrics_dict['target_paddings'],
+    )
 
     return metrics_bundle