code: def test_temporal_localization_model_get_metrics_fn(
      self,
      cls_loss_type,
      box_loss_type,
      output_per_class_displacements,
      expected_box_loss,
      expected_mean_iou,
  ):
      config = ml_collections.ConfigDict(
          {
              "classification_loss_type": cls_loss_type,
              "box_loss_type": box_loss_type,
              "output_per_class_displacements": output_per_class_displacements,
          }
      )
      model = MockTemporalLocalizationModel(config)
      metrics_fn = jax.pmap(model.get_metrics_fn(), axis_name="batch")
      class_logits = np.array(
          [
              [[[1.2], [-0.9], [0.4]], [[-0.4], [-0.8], [-0.1]]],
              [[[1.2], [0.9], [0.4]], [[0.4], [0.8], [0.1]]],
          ]
      )  # shape is (2, 2, 3, 1).
      if output_per_class_displacements:
          pred_displacements = np.array(
              [
                  [
                      [[0.4, 0.4], [0.0, 0.1], [0.2, 0.0]],
                      [[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]],
                  ],
                  [
                      [[1.2, 0.9], [0.4, 0.4], [0.4, 0.2]],
                      [[0.4, 0.8], [0.1, 0.3], [0.3, 0.3]],
                  ],
              ]
          )  # shape is (2, 2, 3, 2).
          logits = np.concatenate(
              [class_logits, pred_displacements], axis=-1
          )
          logits = logits.reshape((2, 2, 9))
      else:
          pred_displacements = np.array(
              [
                  [[0.4, 0.4], [0.1, 0.2]],
                  [[0.4, 0.4], [0.4, 0.8]],
              ]
          )  # shape is (2, 2, 2).
          logits = np.concatenate(
              [np.squeeze(class_logits, axis=-1), pred_displacements],
              axis=-1,
          )
      batch = {
          "batch_mask": np.ones((2,), dtype=np.int32),
          "inputs": {
              "input_mask": np.ones((2, 2), dtype=np.int32),
          },
          "label": np.array(
              [
                  [[1, 0, 0], [0, 0, 0]],
                  [[0, 1, 0], [0, 0, 0]],
              ],
              dtype=np.int32,
          ),
          "displacements": np.array(
              [
                  [
                      [[0.5, 0.4], [0, 0], [0, 0]],
                      [[0, 0], [0, 0], [0, 0]],
                  ],
                  [
                      [[0, 0], [0.3, 0.3], [0, 0]],
                      [[0, 0], [0, 0], [0, 0]],
                  ],
              ],
              dtype=np.float32,
          ),
      }
      if output_per_class_displacements:
          batch["displacements"] = np.array(
              [
                  [
                      [[0.5, 0.4], [0, 0], [0, 0]],
                      [[0, 0], [0, 0], [0, 0]],
                  ],
                  [
                      [[0, 0], [0.3, 0.3], [0, 0]],
                      [[0, 0], [0, 0], [0, 0]],
                  ],
              ],
              dtype=np.float32,
          )
      else:
          batch["displacements"] = np.array(
              [
                  [[0.5, 0.4], [0, 0]],
                  [[0.3, 0.3], [0, 0]],
              ],
              dtype=np.float32,
          )  # shape (2, 2, 2)
      logits, batch = jax_utils.replicate((logits, 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()),
          {"precision@1", 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], 2)
      for key in expected_box_loss_keys:
          self.assertAlmostEqual(
              metrics[key][0], expected_box_loss[key], delta=1e-4
          )
          self.assertAlmostEqual(metrics[key][1], 2)
      self.assertAlmostEqual(metrics["precision@1"][0], 2.0)
      self.assertAlmostEqual(metrics["precision@1"][1], 4)
      self.assertGreaterEqual(metrics[expected_cls_loss_key][0], 0)
      self.assertAlmostEqual(metrics[expected_cls_loss_key][1], 4)
  
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 temporal_localization_base_model
label_desc: Use image node_id 7 for calling the TemporalLocalizationBaseModelTest obj's underlying member method code with example usage: obj.test_temporal_localization_model_get_metrics_fn(cls_loss_type, box_loss_type, output_per_class_displacements, expected_box_loss, expected_mean_iou) without return types