class_docstr: Base model for object detection with matching.
class_docstr_tok: ["Base","model","for","object","detection","with","matching","."]
docstr: Classification loss.
  
  Args:
    outputs: Model predictions. For the purpose of this loss, outputs must
      have key 'pred_logits'. outputs['pred_logits'] is a nd-array of the
      predicted logits of shape [batch-size, num-objects, num-classes].
    batch: Dict that has 'inputs', 'batch_mask' and, 'label' (ground truth).
      batch['label'] is a dict. For the purpose of this loss, label dict must
      have key 'labels', which the value is an int nd-array of labels with
      shape [batch_size, num_boxes, num_classes + 1]. Since the number of
      boxes (objects) in each example in the batch could be different, the
      input pipeline might add padding boxes to some examples. These padding
      boxes are identified based on their class labels. So if the class label
      is `0`, i.e., a one-hot vector of [1, 0, 0, ..., 0], the box/object is a
      padding object and the loss computation will take that into account. The
      input pipeline also pads the partial batches (last batch of eval/test
      set with num_example < batch_size). batch['batch_mask'] is used to
      identify padding examples which is incorporated to set the weight of
      these examples to zero in the loss computations.
    indices: Matcher output of shape [batch-size, 2, num-objects] which
      conveys source to target pairing of objects.
    log: If true, return classification accuracy as well.
  
  Returns:
    loss: Dict with 'loss_class' and other model specific losses.
    metrics: Dict with 'loss_class' and other model specific metrics.
docstr_tok: ["Classification","loss",".","Args",":","outputs",":","Model","predictions",".","For","the","purpose","of","this","loss",",","outputs","must","have","key","'pred_logits","'",".","outputs","[","'pred_logits","'","]","is","a","nd-array","of","the","predicted","logits","of","shape","[","batch-size",",","num-objects",",","num-classes","]",".","batch",":","Dict","that","has","'inputs","'",",","'batch_mask","'","and",",","'label","'","(","ground","truth",")",".","batch","[","'label","'","]","is","a","dict",".","For","the","purpose","of","this","loss",",","label","dict","must","have","key","'labels","'",",","which","the","value","is","an","int","nd-array","of","labels","with","shape","[","batch_size",",","num_boxes",",","num_classes","+","1","]",".","Since","the","number","of","boxes","(","objects",")","in","each","example","in","the","batch","could","be","different",",","the","input","pipeline","might","add","padding","boxes","to","some","examples",".","These","padding","boxes","are","identified","based","on","their","class","labels",".","So","if","the","class","label","is","`","0","`",",","i.e.",",","a","one-hot","vector","of","[","1",",","0",",","0",",","...",",","0","]",",","the","box\/object","is","a","padding","object","and","the","loss","computation","will","take","that","into","account",".","The","input","pipeline","also","pads","the","partial","batches","(","last","batch","of","eval\/test","set","with","num_example","<","batch_size",")",".","batch","[","'batch_mask","'","]","is","used","to","identify","padding","examples","which","is","incorporated","to","set","the","weight","of","these","examples","to","zero","in","the","loss","computations",".","indices",":","Matcher","output","of","shape","[","batch-size",",","2",",","num-objects","]","which","conveys","source","to","target","pairing","of","objects",".","log",":","If","true",",","return","classification","accuracy","as","well",".","Returns",":","loss",":","Dict","with","'loss_class","'","and","other","model","specific","losses",".","metrics",":","Dict","with","'loss_class","'","and","other","model","specific","metrics","."]
code: def labels_losses_and_metrics(
      self,
      outputs: ArrayDict,
      batch: ArrayDict,
      indices: jnp.ndarray,
      log: bool = True,
  ) -> Tuple[ArrayDict, MetricsDict]:
      """Classification loss.
  
      Args:
        outputs: Model predictions. For the purpose of this loss, outputs must
          have key 'pred_logits'. outputs['pred_logits'] is a nd-array of the
          predicted logits of shape [batch-size, num-objects, num-classes].
        batch: Dict that has 'inputs', 'batch_mask' and, 'label' (ground truth).
          batch['label'] is a dict. For the purpose of this loss, label dict must
          have key 'labels', which the value is an int nd-array of labels with
          shape [batch_size, num_boxes, num_classes + 1]. Since the number of
          boxes (objects) in each example in the batch could be different, the
          input pipeline might add padding boxes to some examples. These padding
          boxes are identified based on their class labels. So if the class label
          is `0`, i.e., a one-hot vector of [1, 0, 0, ..., 0], the box/object is a
          padding object and the loss computation will take that into account. The
          input pipeline also pads the partial batches (last batch of eval/test
          set with num_example < batch_size). batch['batch_mask'] is used to
          identify padding examples which is incorporated to set the weight of
          these examples to zero in the loss computations.
        indices: Matcher output of shape [batch-size, 2, num-objects] which
          conveys source to target pairing of objects.
        log: If true, return classification accuracy as well.
  
      Returns:
        loss: Dict with 'loss_class' and other model specific losses.
        metrics: Dict with 'loss_class' and other model specific metrics.
      """
      assert "pred_logits" in outputs
      assert "label" in batch
  
      batch_weights = batch.get("batch_mask")
      losses, metrics = {}, {}
      targets = batch["label"]
      if isinstance(targets, dict):
          targets = targets["labels"]
  
      src_logits = outputs["pred_logits"]
  
      # Apply the permutation communicated by indices.
      src_logits = model_utils.simple_gather(src_logits, indices[:, 0])
      tgt_labels = model_utils.simple_gather(targets, indices[:, 1])
  
      (
          unnormalized_loss_class,
          denom,
      ) = self._compute_per_example_class_loss(
          tgt_labels=tgt_labels,
          src_logits=src_logits,
          batch_weights=batch_weights,
      )
  
      metrics["loss_class"] = (
          unnormalized_loss_class.sum(),
          denom.sum(),
      )
  
      if self.config.normalization == "global":
          denom = jax.lax.pmean(denom.sum(), axis_name="batch")
          denom = jnp.maximum(denom, 1.0)
          normalized_loss_class = unnormalized_loss_class.sum() / denom
      elif self.config.normalization == "per_example":
          normalized_loss_class = unnormalized_loss_class.sum(axis=1)
          denom = jnp.maximum(denom, 1.0)
          normalized_loss_class = (normalized_loss_class / denom).mean()
      else:
          raise ValueError(
              f"Unknown normalization {self.config.normalization}."
          )
  
      losses["loss_class"] = normalized_loss_class
  
      if log:
          # Class accuracy for non-padded (label != 0) labels
          not_padded = tgt_labels[:, :, 0] == 0
          if batch_weights is not None:
              not_padded = not_padded * jnp.expand_dims(
                  batch_weights, axis=1
              )
          num_correct_no_pad = (
              model_utils.weighted_correctly_classified(
                  src_logits[..., 1:],
                  tgt_labels[..., 1:],
                  weights=not_padded,
              )
          )
          metrics["class_accuracy_not_pad"] = (
              num_correct_no_pad,
              not_padded.sum(),
          )
  
      # Sum metrics and normalizers over all replicas.
      for k, v in metrics.items():
          metrics[k] = model_utils.psum_metric_normalizer(v)
      return losses, metrics
  
imports: import abc
  import functools
  from typing import Any, Callable, Dict, Optional, Tuple
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib import matchers
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import box_utils
  from scenic.model_lib.base_models import model_utils
  from scenic.projects.owl_vit import losses
label_desc: Use image node_id 5 for calling the BaseModelWithMatching obj's underlying member method code with example usage: obj.labels_losses_and_metrics(outputs, batch, indices, log) and returns: losses,  metrics