reformatting

priyakasimbeg · priyakasimbeg · commit f1db3d3e6a7f · 2025-03-18T19:27:51.000Z
diff --git a/algoperf/profiler.py b/algoperf/profiler.py
@@ -72,8 +72,8 @@ def _make_report(
                float(np.std(d)),
                len(d),
                float(np.sum(d)),
-               100.0 * float(np.sum(d)) / total_duration)
-              for a, d in self.recorded_durations.items()]
+               100.0 * float(np.sum(d)) / total_duration) for a,
+              d in self.recorded_durations.items()]
     report.sort(key=lambda x: x[5], reverse=True)
     total_calls = sum(x[3] for x in report)
     return report, total_calls, total_duration
diff --git a/algoperf/workloads/criteo1tb/criteo1tb_jax/workload.py b/algoperf/workloads/criteo1tb/criteo1tb_jax/workload.py
@@ -130,16 +130,16 @@ def model_fn(
     return logits_batch, None
 
   @functools.partial(
-    jax.jit,
-    in_shardings=(sharding_utils.get_replicated_sharding(),
-                  sharding_utils.get_naive_sharding_spec(),
-                  ),
-    static_argnums=(0,),
-    out_shardings=sharding_utils.get_replicated_sharding()
-  )
+      jax.jit,
+      in_shardings=(
+          sharding_utils.get_replicated_sharding(),
+          sharding_utils.get_naive_sharding_spec(),
+      ),
+      static_argnums=(0,),
+      out_shardings=sharding_utils.get_replicated_sharding())
   def _eval_batch_jitted(self,
-                          params: spec.ParameterContainer,
-                          batch: Dict[str, spec.Tensor]) -> spec.Tensor:
+                         params: spec.ParameterContainer,
+                         batch: Dict[str, spec.Tensor]) -> spec.Tensor:
     logits, _ = self.model_fn(
         params,
         batch,
@@ -160,8 +160,7 @@ def _eval_batch(self,
                   batch: Dict[str, spec.Tensor]) -> spec.Tensor:
     # We do NOT psum inside of _eval_batch_pmapped, so the returned tensor of
     # shape (local_device_count,) will all be different values.
-    return np.array(
-        self._eval_batch_jitted(params, batch), dtype=np.float64)
+    return np.array(self._eval_batch_jitted(params, batch), dtype=np.float64)
 
 
 class Criteo1TbDlrmSmallTestWorkload(Criteo1TbDlrmSmallWorkload):
diff --git a/algoperf/workloads/fastmri/fastmri_jax/workload.py b/algoperf/workloads/fastmri/fastmri_jax/workload.py
@@ -96,14 +96,11 @@ def loss_fn(
 
   @functools.partial(
       jax.jit,
-      in_shardings=(
-        sharding_utils.get_replicated_sharding(),
-        sharding_utils.get_naive_sharding_spec(),
-        sharding_utils.get_replicated_sharding()
-      ),
+      in_shardings=(sharding_utils.get_replicated_sharding(),
+                    sharding_utils.get_naive_sharding_spec(),
+                    sharding_utils.get_replicated_sharding()),
       static_argnums=(0,),
-      out_shardings=sharding_utils.get_replicated_sharding()
-  )
+      out_shardings=sharding_utils.get_replicated_sharding())
   def _eval_model(self,
                   params: spec.Tensor,
                   batch: Dict[str, spec.Tensor],
diff --git a/algoperf/workloads/fastmri/fastmri_pytorch/workload.py b/algoperf/workloads/fastmri/fastmri_pytorch/workload.py
@@ -250,7 +250,9 @@ def _eval_model_on_split(self,
     for _ in range(num_batches):
       batch = next(self._eval_iters[split])
       batch_metrics = self._eval_model(params, batch, model_rng)
-      total_metrics = {k: v + batch_metrics[k] for k, v in total_metrics.items()}
+      total_metrics = {
+          k: v + batch_metrics[k] for k, v in total_metrics.items()
+      }
     if USE_PYTORCH_DDP:
       for metric in total_metrics.values():
         dist.all_reduce(metric)
diff --git a/algoperf/workloads/imagenet_resnet/imagenet_jax/custom_tf_addons.py b/algoperf/workloads/imagenet_resnet/imagenet_jax/custom_tf_addons.py
@@ -20,31 +20,27 @@
     tf.dtypes.float64,
 }
 
-Number = Union[
-    float,
-    int,
-    np.float16,
-    np.float32,
-    np.float64,
-    np.int8,
-    np.int16,
-    np.int32,
-    np.int64,
-    np.uint8,
-    np.uint16,
-    np.uint32,
-    np.uint64,
-]
-
-TensorLike = Union[
-    List[Union[Number, list]],
-    tuple,
-    Number,
-    np.ndarray,
-    tf.Tensor,
-    tf.SparseTensor,
-    tf.Variable,
-]
+Number = Union[float,
+               int,
+               np.float16,
+               np.float32,
+               np.float64,
+               np.int8,
+               np.int16,
+               np.int32,
+               np.int64,
+               np.uint8,
+               np.uint16,
+               np.uint32,
+               np.uint64,]
+
+TensorLike = Union[List[Union[Number, list]],
+                   tuple,
+                   Number,
+                   np.ndarray,
+                   tf.Tensor,
+                   tf.SparseTensor,
+                   tf.Variable,]
 
 
 def get_ndims(image):
diff --git a/algoperf/workloads/imagenet_resnet/imagenet_jax/randaugment.py b/algoperf/workloads/imagenet_resnet/imagenet_jax/randaugment.py
@@ -316,7 +316,8 @@ def build_lut(histo, step):
     # If step is zero, return the original image.  Otherwise, build
     # lut from the full histogram and step and then index from it.
     result = tf.cond(
-        tf.equal(step, 0), lambda: im,
+        tf.equal(step, 0),
+        lambda: im,
         lambda: tf.gather(build_lut(histo, step), im))
 
     return tf.cast(result, tf.uint8)
@@ -551,6 +552,7 @@ def distort_image_with_randaugment(image, num_layers, magnitude, key):
                                            translate_const=100)
         image = tf.cond(
             tf.equal(i, op_to_select),
-            lambda selected_func=func, selected_args=args: selected_func(
-                image, *selected_args), lambda: image)
+            lambda selected_func=func,
+            selected_args=args: selected_func(image, *selected_args),
+            lambda: image)
   return image
diff --git a/algoperf/workloads/imagenet_resnet/imagenet_jax/workload.py b/algoperf/workloads/imagenet_resnet/imagenet_jax/workload.py
@@ -105,12 +105,12 @@ def init_model_fn(
     self._param_types = param_utils.jax_param_types(self._param_shapes)
     mesh = sharding_utils.get_mesh()
     params = jax.tree_map(
-        lambda x: jax.device_put(x, sharding_utils.get_replicated_sharding(mesh)
-                                ),
+        lambda x: jax.device_put(x,
+                                 sharding_utils.get_replicated_sharding(mesh)),
         params)
     model_state = jax.tree_map(
-        lambda x: jax.device_put(x, sharding_utils.get_replicated_sharding(mesh)
-                                ),
+        lambda x: jax.device_put(x,
+                                 sharding_utils.get_replicated_sharding(mesh)),
         model_state)
     return params, model_state
 
diff --git a/algoperf/workloads/imagenet_resnet/imagenet_pytorch/workload.py b/algoperf/workloads/imagenet_resnet/imagenet_pytorch/workload.py
@@ -307,7 +307,9 @@ def _eval_model_on_split(self,
           update_batch_norm=False)
       weights = batch.get('weights')
       batch_metrics = self._compute_metrics(logits, batch['targets'], weights)
-      total_metrics = {k: v + batch_metrics[k] for k, v in total_metrics.items()}
+      total_metrics = {
+          k: v + batch_metrics[k] for k, v in total_metrics.items()
+      }
     if USE_PYTORCH_DDP:
       for metric in total_metrics.values():
         dist.all_reduce(metric)
diff --git a/algoperf/workloads/librispeech_conformer/librispeech_jax/models.py b/algoperf/workloads/librispeech_conformer/librispeech_jax/models.py
@@ -153,8 +153,8 @@ def setup(self):
     self.kernel = self.param('kernel',
                              nn.initializers.xavier_uniform(),
                              self.filter_shape)
-    self.bias = self.param('bias', lambda rng, s: jnp.zeros(s, jnp.float32),
-                           self.output_channels)
+    self.bias = self.param(
+        'bias', lambda rng, s: jnp.zeros(s, jnp.float32), self.output_channels)
 
   @nn.compact
   def __call__(self, inputs, paddings):
@@ -442,10 +442,12 @@ def setup(self):
     dtype = self.config.dtype
 
     self.ra_mean = self.variable('batch_stats',
-                                 'mean', lambda s: jnp.zeros(s, dtype),
+                                 'mean',
+                                 lambda s: jnp.zeros(s, dtype),
                                  dim)
     self.ra_var = self.variable('batch_stats',
-                                'var', lambda s: jnp.ones(s, dtype),
+                                'var',
+                                lambda s: jnp.ones(s, dtype),
                                 dim)
 
     self.gamma = self.param('scale', nn.initializers.zeros, dim, dtype)
diff --git a/algoperf/workloads/librispeech_conformer/librispeech_jax/spectrum_augmenter.py b/algoperf/workloads/librispeech_conformer/librispeech_jax/spectrum_augmenter.py
@@ -81,8 +81,8 @@ def _get_mask(self,
           jnp.expand_dims(jnp.arange(multiplicity, dtype=jnp.int32), 0),
           [batch_size, 1])
       multiplicity_tensor = masks_per_frame * choose_range
-      multiplicity_weights = (multiplicity_weights
-                              < multiplicity_tensor).astype(jnp.int32)
+      multiplicity_weights = (multiplicity_weights <
+                              multiplicity_tensor).astype(jnp.int32)
       pre_mask = jnp.einsum('bmt,bm->bt', pre_mask, multiplicity_weights)
     else:
       pre_mask = jnp.einsum('bmt->bt', pre_mask)
diff --git a/algoperf/workloads/librispeech_conformer/librispeech_jax/workload.py b/algoperf/workloads/librispeech_conformer/librispeech_jax/workload.py
@@ -446,7 +446,7 @@ def use_gelu(self) -> bool:
   @property
   def validation_target_value(self) -> float:
     return 0.094114
- 
+
   @property
   def test_target_value(self) -> float:
     return 0.056629
diff --git a/algoperf/workloads/librispeech_conformer/librispeech_pytorch/workload.py b/algoperf/workloads/librispeech_conformer/librispeech_pytorch/workload.py
@@ -259,9 +259,8 @@ def greedy_decode(
     idxs = torch.arange(
         fin_result.numel(), device=result.device).view(*fin_result.shape)
     mask = torch.arange(
-        fin_result.shape[1],
-        device=result.device).view(1, -1) < result.count_nonzero(dim=1).view(
-            -1, 1)
+        fin_result.shape[1], device=result.device).view(
+            1, -1) < result.count_nonzero(dim=1).view(-1, 1)
     fin_result.view(-1)[idxs[mask != 0]] = result[result != blank_id]
     padding = fin_result == 0
     return fin_result, padding
@@ -329,7 +328,9 @@ def _eval_model_on_split(self,
           'word_errors': word_errors,
           'num_words': num_words,
       }
-      total_metrics = {k: v + batch_metrics[k] for k, v in total_metrics.items()}
+      total_metrics = {
+          k: v + batch_metrics[k] for k, v in total_metrics.items()
+      }
     if USE_PYTORCH_DDP:
       for metric in total_metrics.values():
         dist.all_reduce(metric)
diff --git a/algoperf/workloads/librispeech_deepspeech/librispeech_jax/models.py b/algoperf/workloads/librispeech_deepspeech/librispeech_jax/models.py
@@ -144,8 +144,8 @@ def setup(self):
     self.kernel = self.param('kernel',
                              nn.initializers.xavier_uniform(),
                              self.filter_shape)
-    self.bias = self.param('bias', lambda rng, s: jnp.zeros(s, jnp.float32),
-                           self.output_channels)
+    self.bias = self.param(
+        'bias', lambda rng, s: jnp.zeros(s, jnp.float32), self.output_channels)
 
   @nn.compact
   def __call__(self, inputs, paddings, train):
@@ -278,10 +278,12 @@ def setup(self):
     dtype = self.dtype
 
     self.ra_mean = self.variable('batch_stats',
-                                 'mean', lambda s: jnp.zeros(s, dtype),
+                                 'mean',
+                                 lambda s: jnp.zeros(s, dtype),
                                  dim)
     self.ra_var = self.variable('batch_stats',
-                                'var', lambda s: jnp.ones(s, dtype),
+                                'var',
+                                lambda s: jnp.ones(s, dtype),
                                 dim)
 
     self.gamma = self.param('scale', nn.initializers.zeros, dim, dtype)
@@ -393,7 +395,7 @@ def __call__(
       c_0_shape = np.shape(c_0)
       weights_shape = np.shape(weights)
       seq_lengths_np = np.shape(seq_lengths)
-      
+
       n = jax.devices()
       logging.info(f"jax num devices {n}")
       logging.info(f'inputs shape {inputs_shape}')
@@ -440,6 +442,7 @@ def unpack_weights(
         self.bidirectional,
     )
 
+
 ### Swap in regular LSTM layer for debuggin
 @jax.vmap
 def flip_sequences(inputs: Array, lengths: Array) -> Array:
@@ -474,6 +477,7 @@ def flip_sequences(inputs: Array, lengths: Array) -> Array:
   idxs = (jnp.arange(max_length - 1, -1, -1) + lengths) % max_length
   return inputs[idxs]
 
+
 class GenericRNNSequenceEncoder(nn.Module):
   """Encodes a single sequence using any RNN cell, for example `nn.LSTMCell`.
 
@@ -503,8 +507,11 @@ def setup(self):
       in_axes=(1, flax.core.axes_scan.broadcast, flax.core.axes_scan.broadcast),
       out_axes=1,
       split_rngs={'params': False})
-  def unroll_cell(self, cell_state: StateType, inputs: Array,
-                  recurrent_dropout_mask: Optional[Array], deterministic: bool):
+  def unroll_cell(self,
+                  cell_state: StateType,
+                  inputs: Array,
+                  recurrent_dropout_mask: Optional[Array],
+                  deterministic: bool):
     """Unrolls a recurrent cell over an input sequence.
 
     Args:
@@ -554,7 +561,8 @@ def __call__(self,
       inputs = flip_sequences(inputs, lengths)
 
     recurrent_dropout_mask = None
-    _, (cell_states, outputs) = self.unroll_cell(initial_state, inputs,
+    _, (cell_states, outputs) = self.unroll_cell(initial_state,
+                                                 inputs,
                                                  recurrent_dropout_mask,
                                                  deterministic)
     final_state = jax.tree.map(
@@ -602,8 +610,7 @@ def __call__(
       inputs: Array,
       lengths: Array,
       initial_states: Optional[Sequence[StateType]] = None,
-      deterministic: bool = False
-  ) -> Tuple[Array, Sequence[StateType]]:
+      deterministic: bool = False) -> Tuple[Array, Sequence[StateType]]:
     """Processes the input sequence using the recurrent cell.
 
     Args:
@@ -635,15 +642,12 @@ def __call__(
       rng = jax.random.PRNGKey(0)
       initial_states = [
           self.cell_type(self.hidden_size).initialize_carry(
-              rng, (batch_size, 1)
-          )
-          for _ in range(num_cells)
+              rng, (batch_size, 1)) for _ in range(num_cells)
       ]
     if len(initial_states) != num_cells:
       raise ValueError(
           f'Please provide {self.num_cells} (`num_layers`, *2 if bidirectional)'
-          'initial states.'
-      )
+          'initial states.')
 
     # For each layer, apply the forward and optionally the backward RNN cell.
     cell_idx = 0
@@ -756,6 +760,7 @@ def __call__(
             initial_states=initial_states,
             deterministic=deterministic)
 
+
 class BatchRNN(nn.Module):
   """Implements a single deepspeech encoder layer.
   """
@@ -775,10 +780,9 @@ def __call__(self, inputs, input_paddings, train):
                                                     input_paddings,
                                                     train)
 
-    # For regular LSTM                                                
+    # For regular LSTM
     hidden_size = (
-      config.encoder_dim // 2 if config.bidirectional else config.encoder_dim
-    )
+        config.encoder_dim // 2 if config.bidirectional else config.encoder_dim)
     lengths = jnp.sum(1 - input_paddings, axis=-1, dtype=jnp.int32)
 
     # output, _ = LSTM(
@@ -791,7 +795,6 @@ def __call__(self, inputs, input_paddings, train):
         features=config.encoder_dim // 2,
         bidirectional=config.bidirectional,
         num_layers=1)(inputs, input_paddings)
-    
 
     return output
 
diff --git a/algoperf/workloads/mnist/workload.py b/algoperf/workloads/mnist/workload.py
@@ -46,7 +46,8 @@ def _build_mnist_dataset(
   ds = ds.map(
       lambda x: {
           'inputs': _normalize(x['image'], train_mean, train_stddev),
-          'targets': x['label'],})
+          'targets': x['label'],
+      })
   is_train = split == 'train'
 
   if cache:
@@ -213,6 +214,8 @@ def _eval_model_on_split(self,
                                        batch,
                                        model_state,
                                        per_device_model_rngs)
-      total_metrics = {k: v + batch_metrics[k] for k, v in total_metrics.items()}
+      total_metrics = {
+          k: v + batch_metrics[k] for k, v in total_metrics.items()
+      }
 
     return self._normalize_eval_metrics(num_examples, total_metrics)
diff --git a/algoperf/workloads/wmt/wmt_pytorch/models.py b/algoperf/workloads/wmt/wmt_pytorch/models.py
