class_docstr: Test DETRModel.
class_docstr_tok: ["Test","DETRModel","."]
docstr: Test loss_labels by checking its output's dictionary format.
  
  Args:
    log: bool; Whether do logging or not in labels_losses_and_metrics.
    metrics_key: list; Expected metric keys.
docstr_tok: ["Test","loss_labels","by","checking","its","output","'s","dictionary","format",".","Args",":","log",":","bool",";","Whether","do","logging","or","not","in","labels_losses_and_metrics",".","metrics_key",":","list",";","Expected","metric","keys","."]
code: def test_labels_losses_and_metrics(self, log, metrics_key):
      """Test loss_labels by checking its output's dictionary format.
  
      Args:
        log: bool; Whether do logging or not in labels_losses_and_metrics.
        metrics_key: list; Expected metric keys.
      """
  
      def f_to_pmap(outputs, batch):
          return self.model.labels_losses_and_metrics(
              outputs, batch, indices=self.indices, log=log
          )
  
      # test loss function in the pmapped setup
      labels_lm_pmapped = jax.pmap(f_to_pmap, axis_name="batch")
  
      outputs, batch = (
          jax_utils.replicate(self.outputs),
          jax_utils.replicate(self.batch),
      )
  
      losses, metrics = labels_lm_pmapped(outputs, batch)
      losses = jax_utils.unreplicate(losses)
      metrics = jax_utils.unreplicate(metrics)
      self.assertSameElements(losses.keys(), ["loss_class"])
      self.is_valid_loss(losses["loss_class"])
      self.assertSameElements(metrics.keys(), metrics_key)
      for mk in metrics_key:
          self.is_valid(metrics[mk][0])
          self.is_valid(metrics[mk][1])
  
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  from flax import jax_utils
  from flax import traverse_util
  from flax.core import unfreeze
  import jax
  from jax import random
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.projects.baselines.detr import model
label_desc: Use image node_id 4 for calling the DETRModelTest obj's underlying member method code with example usage: obj.test_labels_losses_and_metrics(log, metrics_key) and returns: self