class_docstr: Tests the default trainer on single device setup.
class_docstr_tok: ["Tests","the","default","trainer","on","single","device","setup","."]
docstr: Test trainer evaluate end to end with classification model metrics.
docstr_tok: ["Test","trainer","evaluate","end","to","end","with","classification","model","metrics","."]
code: def test_classifaction_model_evaluate(self, get_confusion_matrix):
      """Test trainer evaluate end to end with classification model metrics."""
      # Define a fixed output for the fake flax model.
      fake_batch_logits = np.tile([0.5, 0.2, 0.7, 0.0], (4, 1))
      # 4 evaluation batches of size 4.
      fake_batches = [
          {
              "inputs": None,
              "label": np.array([3, 2, 1, 0]),
              "batch_mask": np.array([1, 1, 1, 1]),
          },
          {
              "inputs": None,
              "label": np.array([0, 3, 2, 0]),
              "batch_mask": np.array([1, 1, 1, 1]),
          },
          {
              "inputs": None,
              "label": np.array([0, 0, 0, 0]),
              "batch_mask": np.array([1, 1, 1, 1]),
          },
          {
              "inputs": None,
              "label": np.array([1, 1, 1, 1]),
              "batch_mask": np.array([1, 1, 1, 1]),
          },
      ]
  
      rng = jax.random.PRNGKey(0)
      model, train_state = self.get_train_state(rng, fake_batch_logits)
      eval_summary = self.train_and_evaluation(
          model,
          train_state,
          fake_batches,
          functools.partial(
              classification_model.classification_metrics_function,
              target_is_onehot=False,
          ),
          get_confusion_matrix,
      )
      if get_confusion_matrix:
          eval_summary, confusion_matrix_summary = eval_summary
  
      def batch_loss(logits, targets):
          # Softmax cross-entropy loss.
          one_hot_targets = np.eye(4)[targets]
          loss = -np.sum(
              one_hot_targets * nn.log_softmax(logits), axis=-1
          )
          return loss
  
      expected_accuracy = (
          2.0 / 16.0
      )  # FakeFlaxModule always predicts class 2.
      expected_loss = np.mean(
          [
              batch_loss(fake_batch_logits, b["label"])
              for b in fake_batches
          ]
      )
  
      self.assertEqual(expected_accuracy, eval_summary["accuracy"])
      np.testing.assert_allclose(
          expected_loss, eval_summary["loss"], atol=1e-6
      )
  
      if get_confusion_matrix:
          # As FakeFlaxModule always predicts class 2, this class has a recall of 1,
          # and the denominator for the precision is the total number of examples.
          self.assertAlmostEqual(
              confusion_matrix_summary["precision/0"], 0.0
          )
          self.assertAlmostEqual(
              confusion_matrix_summary["precision/1"], 0.0
          )
          self.assertAlmostEqual(
              confusion_matrix_summary["precision/2"], 2.0 / 16.0
          )
          self.assertAlmostEqual(
              confusion_matrix_summary["precision/3"], 0.0
          )
          self.assertAlmostEqual(
              confusion_matrix_summary["recall/0"], 0.0
          )
          self.assertAlmostEqual(
              confusion_matrix_summary["recall/1"], 0.0
          )
          self.assertAlmostEqual(
              confusion_matrix_summary["recall/2"], 1.0
          )
          self.assertAlmostEqual(
              confusion_matrix_summary["recall/3"], 0.0
          )
          self.assertAlmostEqual(
              confusion_matrix_summary["jaccard/0"], 0.0
          )
          self.assertAlmostEqual(
              confusion_matrix_summary["jaccard/1"], 0.0
          )
          self.assertAlmostEqual(
              confusion_matrix_summary["jaccard/2"], 2.0 / 16.0
          )
          self.assertAlmostEqual(
              confusion_matrix_summary["jaccard/3"], 0.0
          )
  
imports: import functools
  import shutil
  import os
  from absl.testing import absltest
  from absl.testing import parameterized
  import flax
  from flax import jax_utils
  import flax.linen
  import jax.numpy
  import jax.random
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.base_models import multilabel_classification_model
  from scenic.projects.vivit import evaluation_lib
  from scenic.projects.vivit import train_utils
  from scenic.train_lib_deprecated import optimizers
  from scenic.train_lib_deprecated import train_utils
  import tensorflow
label_desc: Use image node_id 5 for calling the ViViTClassificationTrainerTest obj's underlying member method code with example usage: obj.test_classifaction_model_evaluate(get_confusion_matrix) and returns: loss