remove _eval_batch from jax workload

Author: priyakasimbeg

algoperf/workloads/lm/lm_jax/workload.py

@@ -139,20 +139,3 @@ def _normalize_eval_metrics(
     del num_examples
     eval_denominator = total_metrics.pop('denominator')
     return jax.tree.map(lambda x: float(x / eval_denominator), total_metrics)
-
-
-  def _eval_batch(self,
-                  params: spec.ParameterContainer,
-                  batch: Dict[str, spec.Tensor],
-                  model_state: spec.ModelAuxiliaryState,
-                  rng: spec.RandomState) -> spec.Tensor:
-    """Evaluate the model on a single batch."""
-    logits, _ = self.model_fn(
-        params, batch, model_state, spec.ForwardPassMode.EVAL, rng, False)
-    # Calculate cross-entropy loss
-    # TODO(kasimbeg): add weights?
-    metrics = self.compute_weighted_cross_entropy(logits, batch['targets'], batch['weights'])
-    return {
-      'loss': metrics['summed'],
-      'denominator': metrics['n_valid_examples'],
-    }

algoperf/workloads/lm/workload.py

@@ -124,25 +124,6 @@ def _build_input_queue(
   ):
     """Build an input queue for the given split."""
 
-  def _eval_batch(
-    self,
-    params: spec.ParameterContainer,
-    batch: Dict[str, spec.Tensor],
-    model_state: spec.ModelAuxiliaryState,
-    rng: spec.RandomState,
-  ) -> spec.Tensor:
-    """Evaluate the model on a single batch."""
-    logits, _ = self.model_fn(
-      params,
-      batch,
-      model_state,
-      spec.ForwardPassMode.EVAL,
-      rng,
-      update_batch_norm=False,
-    )
-
-    loss_dict = self.loss_fn(batch['targets'], logits)
-    return loss_dict
 
   def _eval_model_on_split(
     self,
@@ -181,6 +162,23 @@ def _eval_model_on_split(
     eval_results['ppl'] = np.exp(eval_results['loss'])      
     return eval_results
 
+
+  def _eval_batch(self,
+                  params: spec.ParameterContainer,
+                  batch: Dict[str, spec.Tensor],
+                  model_state: spec.ModelAuxiliaryState,
+                  rng: spec.RandomState) -> spec.Tensor:
+    """Evaluate the model on a single batch."""
+    logits, _ = self.model_fn(
+        params, batch, model_state, spec.ForwardPassMode.EVAL, rng, False)
+    # Calculate cross-entropy loss
+    metrics = self.compute_weighted_cross_entropy(logits, batch['targets'], batch['weights'])
+    return {
+      'loss': metrics['summed'],
+      'denominator': metrics['n_valid_examples'],
+    }
+
+
   @abc.abstractmethod
   def _normalize_eval_metrics(
     self, num_examples: int, total_metrics: Dict[str, Any]