class_docstr: Base model for object detection or instance segmentation with matching.
  
  This model implements the classification and matching parts which are shared
  between the object detection and instance segmentation models.
docstr: Classification softmax cross entropy loss and (optionally) top-1 correct.
  
  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 pre-softmax 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. It may
      be one-hot if dataset_meta_data.target_is_onehot was set to True. If
      batch['batch_mask'] is provided it is used to weight the loss for
      different images in the current batch of examples.
    indices: Matcher output of shape [batch-size, 2, num-objects] which
      conveys source to target pairing of objects.
    log: If true, return class_error as well.
  
  Returns:
    loss: Dict with keys 'loss_class'.
    metrics: Dict with keys 'loss_class' and 'class_error`.
docstr_tok: ["Classification","softmax","cross","entropy","loss","and","(","optionally",")","top-1","correct",".","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","pre-softmax","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",".","It","may","be","one-hot","if","dataset_meta_data.target_is_onehot","was","set","to","True",".","If","batch","[","'batch_mask","'","]","is","provided","it","is","used","to","weight","the","loss","for","different","images","in","the","current","batch","of","examples",".","indices",":","Matcher","output","of","shape","[","batch-size",",","2",",","num-objects","]","which","conveys","source","to","target","pairing","of","objects",".","log",":","If","true",",","return","class_error","as","well",".","Returns",":","loss",":","Dict","with","keys","'loss_class","'",".","metrics",":","Dict","with","keys","'loss_class","'","and","'class_error","`","."]
code: def labels_losses_and_metrics(
      self,
      outputs: ArrayDict,
      batch: ArrayDict,
      indices: jnp.ndarray,
      log: bool = True,
  ) -> Tuple[ArrayDict, MetricsDict]:
      """Classification softmax cross entropy loss and (optionally) top-1 correct.
  
      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 pre-softmax 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. It may
          be one-hot if dataset_meta_data.target_is_onehot was set to True. If
          batch['batch_mask'] is provided it is used to weight the loss for
          different images in the current batch of examples.
        indices: Matcher output of shape [batch-size, 2, num-objects] which
          conveys source to target pairing of objects.
        log: If true, return class_error as well.
  
      Returns:
        loss: Dict with keys 'loss_class'.
        metrics: Dict with keys 'loss_class' and 'class_error`.
      """
      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"]
      num_classes = src_logits.shape[-1]
  
      # Copy unpermuted logits & targets; this is done to be doubly sure that
      # prec@1 evaluation is not in any way affected by the matching.
      orig_src_logits, orig_tgt_labels = src_logits, targets
  
      # 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])
      if self.dataset_meta_data.get("target_is_onehot", False):
          tgt_labels_onehot = tgt_labels
          orig_tgt_labels_onehot = orig_tgt_labels
      else:
          tgt_labels_onehot = jax.nn.one_hot(tgt_labels, num_classes)
          orig_tgt_labels_onehot = jax.nn.one_hot(
              orig_tgt_labels, num_classes
          )
  
      def get_label(label):
          if (
              isinstance(batch["label"], dict)
              and label in batch["label"]
          ):
              value = batch["label"][label]
          else:
              value = None
          return value
  
      neg_category_ids = get_label("neg_category_ids")
      loose_box = get_label("loose_box")
  
      # Normalization before masking is important so that masked classes can
      # be assigned when using softmax normalization. In federated datasets
      # we don't have exhaustive annotation so this is important.
      src_log_p = self.logits_to_probs(src_logits, log_p=True)
  
      (
          unnormalized_loss_class,
          denom,
      ) = self._compute_per_example_class_loss(
          tgt_labels_onehot=tgt_labels_onehot,
          src_log_p=src_log_p,
          batch_weights=batch_weights,
          neg_category_ids=neg_category_ids,
          loose_box=loose_box,
      )
  
      metrics["loss_class"] = (
          unnormalized_loss_class.sum(),
          denom.sum(),
      )
  
      norm_type = self.config.get("normalization", "detr")
      if norm_type == "detr":
          denom = jnp.maximum(denom.sum(), 1.0)
          normalized_loss_class = unnormalized_loss_class.sum() / denom
      elif norm_type == "global":
          denom = jax.lax.pmean(denom.sum(), axis_name="batch")
          denom = jnp.maximum(denom, 1.0)
          normalized_loss_class = unnormalized_loss_class.sum() / denom
      else:
          raise ValueError(f"Unknown normalization {norm_type}.")
  
      losses["loss_class"] = normalized_loss_class
  
      if log:
          # For normalization, we need to have number of inputs that we do
          # prediction for, which is number of examples in the batch times
          # number of boxes (including padded boxes).
          # note that: tgt_labels_onehot.shape = (bs, num_boxes, num_classes)
          if batch_weights is not None:
              batch_num_inputs = (
                  batch_weights.sum() * tgt_labels_onehot.shape[-2]
              )
          else:
              batch_num_inputs = np.prod(tgt_labels_onehot.shape[:-1])
  
          # Class accuracy for non-padded (label != 0) labels
          not_padded = tgt_labels_onehot[:, :, 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_log_p[..., 1:],
                  tgt_labels_onehot[..., 1:],
                  weights=not_padded,
              )
          )
          metrics["class_accuracy_not_pad"] = (
              num_correct_no_pad,
              not_padded.sum(),
          )
  
          if not self.config.get("sigmoid_loss", False):
              num_correct = model_utils.weighted_correctly_classified(
                  src_log_p, tgt_labels_onehot, weights=batch_weights
              )
              metrics["class_accuracy"] = (
                  num_correct,
                  batch_num_inputs,
              )
  
          if "loss_bbox" not in self.loss_terms_weights:
              # Prec@1 for classification only models.
              if batch_weights is not None:
                  batch_num_inputs = batch_weights.sum()
              else:
                  batch_num_inputs = jnp.asarray(
                      tgt_labels_onehot.shape[0]
                  )
              max_logits = jnp.max(orig_src_logits, axis=-2)
              tgt_labels_multihot = jnp.max(
                  orig_tgt_labels_onehot, axis=-2
              )
              prec_at_one = (
                  model_utils.weighted_top_one_correctly_classified(
                      max_logits[:, 1:],
                      tgt_labels_multihot[:, 1:],
                      weights=batch_weights,
                  )
              )
              metrics["prec@1"] = (prec_at_one, batch_num_inputs)
  
      # 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
  import numpy
  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
label_desc: Use image node_id 6 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, value