fixing lint except yapf

.github/workflows/CI.yml

@@ -21,6 +21,7 @@ jobs:
     - name: Run pytest
       run: |
         pytest -vx tests/version_test.py
+        pytest -vx tests/reference_submission_tests.py
         pytest -vx tests/workloads/imagenet/imagenet_jax/workload_test.py
         pytest -vx tests/test_num_params.py
         pytest -vx tests/test_param_shapes.py

.github/workflows/linting.yml

@@ -19,6 +19,7 @@ jobs:
       run: |
         pylint algorithmic_efficiency
         pylint baselines
+        pylint reference_submissions
         pylint submission_runner.py
         pylint tests
 

algorithmic_efficiency/spec.py

@@ -3,8 +3,10 @@
 import abc
 import enum
 from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union
+
 from absl import logging
 
+
 class LossType(enum.Enum):
   SOFTMAX_CROSS_ENTROPY = 0
   SIGMOID_CROSS_ENTROPY = 1

algorithmic_efficiency/workloads/imagenet/imagenet_pytorch/workload.py

@@ -1,8 +1,8 @@
 """ImageNet workload implemented in PyTorch."""
 
 import contextlib
-import os
 import math
+import os
 from typing import Tuple
 
 import torch

algorithmic_efficiency/workloads/librispeech/librispeech_pytorch/workload.py

@@ -16,7 +16,6 @@
 from algorithmic_efficiency.workloads.librispeech.librispeech_pytorch import \
     models
 
-
 device = torch.device(
         "cuda:0" if torch.cuda.is_available() else "cpu")
 

algorithmic_efficiency/workloads/wmt/wmt_jax/workload.py

@@ -11,7 +11,6 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-import tensorflow as tf
 
 from algorithmic_efficiency import spec
 from algorithmic_efficiency.workloads.wmt import bleu
@@ -89,7 +88,10 @@ def compute_weighted_cross_entropy(self,
         ((self._vocab_size - 1) * low_confidence *
          jnp.log(low_confidence + 1e-20)))
     soft_targets = common_utils.onehot(
-        targets, self._vocab_size, on_value=confidence, off_value=low_confidence)
+        targets,
+        self._vocab_size,
+        on_value=confidence,
+        off_value=low_confidence)
 
     loss = -jnp.sum(soft_targets * nn.log_softmax(logits), axis=-1)
     loss = loss - normalizing_constant

reference_submissions/imagenet/imagenet_jax/submission.py

@@ -24,7 +24,8 @@ def cosine_decay(lr, step, total_steps):
   return mult * lr
 
 
-def create_learning_rate_fn(hparams: spec.Hyperparameters, steps_per_epoch: int):
+def create_learning_rate_fn(
+    hparams: spec.Hyperparameters, steps_per_epoch: int):
   """Create learning rate schedule."""
   base_learning_rate = hparams.learning_rate * get_batch_size('imagenet') / 256.
   warmup_fn = optax.linear_schedule(

reference_submissions/mnist/mnist_pytorch/submission.py

@@ -64,7 +64,8 @@ def update_params(
       rng=rng,
       update_batch_norm=True)
 
-  loss = workload.loss_fn(label_batch=batch['targets'], logits_batch=output).mean()
+  loss = workload.loss_fn(
+      label_batch=batch['targets'], logits_batch=output).mean()
 
   loss.backward()
   optimizer_state['optimizer'].step()

reference_submissions/wmt/wmt_jax/submission.py

@@ -5,11 +5,9 @@
 
 from flax import jax_utils
 from flax import linen as nn
-from flax import optim
 from flax.training import common_utils
 import jax
 import jax.numpy as jnp
-import numpy as np
 import optax
 
 from algorithmic_efficiency import spec

tests/reference_submission_tests.py

@@ -220,18 +220,18 @@ def _test_submission(workload_name,
       hyperparameters,
       global_step,
       data_select_rng)
-  # _, model_params, model_state = update_params(
-  #     workload=workload,
-  #     current_param_container=model_params,
-  #     current_params_types=workload.model_params_types,
-  #     model_state=model_state,
-  #     hyperparameters=hyperparameters,
-  #     batch=batch,
-  #     loss_type=workload.loss_type,
-  #     optimizer_state=optimizer_state,
-  #     eval_results=[],
-  #     global_step=global_step,
-  #     rng=update_rng)
+  _, model_params, model_state = update_params(
+      workload=workload,
+      current_param_container=model_params,
+      current_params_types=workload.model_params_types,
+      model_state=model_state,
+      hyperparameters=hyperparameters,
+      batch=batch,
+      loss_type=workload.loss_type,
+      optimizer_state=optimizer_state,
+      eval_results=[],
+      global_step=global_step,
+      rng=update_rng)
   eval_result = workload.eval_model(global_batch_size,
                                     model_params,
                                     model_state,
@@ -261,11 +261,6 @@ def test_submission(self):
       for framework in ['jax', 'pytorch']:
         submission_dir = f'{workload_dir}/{workload_name}_{framework}'
         if os.path.exists(submission_dir):
-          # # DO NOT SUBMIT
-          # if 'mnist' in submission_dir or 'imagenet' in submission_dir or 'librispeech' in submission_dir or 'ogbg' in submission_dir:
-          #   continue
-          # if not ('librispeech' in submission_dir):
-          #   continue
           submission_path = (f'reference_submissions/{workload_name}/'
                              f'{workload_name}_{framework}/submission.py')
           logging.info(f'========= Testing {workload_name} in {framework}.')