code: def test_moment_retrieval_model_get_metrics_fn(
      self,
      cls_loss_type,
      box_loss_type,
      all_gather_loss,
      expected_box_loss,
      expected_mean_iou,
  ):
      config = ml_collections.ConfigDict(
          {
              "classification_loss_type": cls_loss_type,
              "box_loss_type": box_loss_type,
              "all_gather_loss": all_gather_loss,
          }
      )
      model = MockMomentRetrievalModel(config)
      metrics_fn = jax.pmap(model.get_metrics_fn(), axis_name="batch")
      logits, batch = jax_utils.replicate((self.logits, self.batch))
      metrics = metrics_fn(logits, batch)
      expected_cls_loss_key = (
          "sigmoid_classification_loss"
          if cls_loss_type == "sigmoid"
          else "focal_classification_loss"
      )
      expected_box_loss_keys = set(expected_box_loss.keys())
      self.assertSetEqual(
          set(metrics.keys()),
          {"accuracy", expected_cls_loss_key, "mean_iou"}
          | expected_box_loss_keys,
      )
      metrics = jax_utils.unreplicate(metrics)
      self.assertAlmostEqual(
          metrics["mean_iou"][0], expected_mean_iou, delta=1e-4
      )
      self.assertAlmostEqual(metrics["mean_iou"][1], 6)
      for key in expected_box_loss_keys:
          self.assertAlmostEqual(
              metrics[key][0], expected_box_loss[key], delta=1e-4
          )
          self.assertAlmostEqual(metrics[key][1], 6)
      if all_gather_loss:
          self.assertAlmostEqual(metrics["accuracy"][0], 16)
      else:
          self.assertAlmostEqual(metrics["accuracy"][0], 8)
  
      if all_gather_loss:
          self.assertAlmostEqual(metrics["accuracy"][1], 16)
      else:
          self.assertAlmostEqual(metrics["accuracy"][1], 8)
  
      self.assertGreaterEqual(metrics[expected_cls_loss_key][0], 0)
  
      if all_gather_loss:
          self.assertGreaterEqual(metrics[expected_cls_loss_key][1], 16)
      else:
          self.assertGreaterEqual(metrics[expected_cls_loss_key][1], 8)
  
file_path: scenic/scenic/projects/unloc/moment_retrieval_base_model_test.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  from flax import jax_utils
  import jax
  import ml_collections
  import numpy
  from scenic.projects.unloc import moment_retrieval_base_model
label_desc: Use image node_id 3 for calling the MomentRetrievalBaseModelTest obj's underlying member method code with example usage: obj.test_moment_retrieval_model_get_metrics_fn(cls_loss_type, box_loss_type, all_gather_loss, expected_box_loss, expected_mean_iou) without return types