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 with multi-label classification model metrics.
docstr_tok: ["Test","trainer","evaluate","with","multi-label","classification","model","metrics","."]
code: def test_multi_label_classifaction_model_evaluate(self):
      """Test trainer evaluate with multi-label classification model metrics."""
      # define a fix 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, with multihot labels.
      fake_batches = [
          {
              "inputs": None,
              "label": np.array(
                  [
                      [1, 0, 0, 1],
                      [0, 0, 1, 0],
                      [0, 1, 0, 0],
                      [1, 0, 0, 1],
                  ]
              ),
          },
          {
              "inputs": None,
              "label": np.array(
                  [
                      [1, 1, 0, 0],
                      [0, 0, 0, 1],
                      [0, 0, 1, 1],
                      [1, 0, 0, 1],
                  ]
              ),
          },
          {
              "inputs": None,
              "label": np.array(
                  [
                      [1, 0, 0, 1],
                      [1, 0, 0, 0],
                      [1, 1, 0, 0],
                      [1, 0, 0, 0],
                  ]
              ),
          },
          {
              "inputs": None,
              "label": np.array(
                  [
                      [0, 1, 0, 1],
                      [0, 1, 0, 0],
                      [1, 1, 0, 0],
                      [0, 1, 0, 0],
                  ]
              ),
          },
      ]
  
      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(
              multilabel_classification_model.multilabel_classification_metrics_function,
              target_is_multihot=True,
          ),
      )
  
      def batch_loss(logits, multi_hot_targets):
          # sigmoid cross-entropy loss
          log_p = jax.nn.log_sigmoid(logits)
          log_not_p = jax.nn.log_sigmoid(-logits)
          loss = -np.sum(
              multi_hot_targets * log_p
              + (1.0 - multi_hot_targets) * log_not_p,
              axis=-1,
          )
          return loss
  
      expected_prec_at_one = (
          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_prec_at_one, eval_summary["prec@1"])
      np.testing.assert_allclose(
          expected_loss, eval_summary["loss"], atol=1e-6
      )
  
imports: import functools
  import shutil
  import tempfile
  from absl.testing import absltest
  from clu import metric_writers
  import flax
  from flax import jax_utils
  import flax.linen
  import jax.numpy
  import jax.random
  import ml_collections
  import numpy
  from scenic.dataset_lib import datasets
  from scenic.model_lib import models
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.base_models import multilabel_classification_model
  from scenic.train_lib import classification_trainer
  from scenic.train_lib import lr_schedules
  from scenic.train_lib import optimizers
  from scenic.train_lib import train_utils
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 6 for calling the ClassificationTrainerTest obj's underlying member method code with example usage: obj.test_multi_label_classifaction_model_evaluate() and returns: loss