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: Loss and metrics function for ObjectDetectionWithMatchingModel.
  
  Args:
    outputs: Model prediction. The exact fields depend on the losses used.
      Please see labels_losses_and_metrics and boxes_losses_and_metrics for
      details.
    batch: Dict that has 'inputs', 'batch_mask' and, 'label' (ground truth).
      batch['label'] is a dict where the keys and values depend on the losses
      used. Please see labels_losses_and_metrics and boxes_losses_and_metrics
      member methods.
    matches: Output of hungarian matcher, if available.
    model_params: pytree (optional); Parameters of the model.
  
  Returns:
    total_loss: Total loss weighted appropriately using
      self.loss_terms_weights and combined across all auxiliary outputs.
    metrics_dict: Individual loss terms with and without weighting for
      logging purposes.
docstr_tok: ["Loss","and","metrics","function","for","ObjectDetectionWithMatchingModel",".","Args",":","outputs",":","Model","prediction",".","The","exact","fields","depend","on","the","losses","used",".","Please","see","labels_losses_and_metrics","and","boxes_losses_and_metrics","for","details",".","batch",":","Dict","that","has","'inputs","'",",","'batch_mask","'","and",",","'label","'","(","ground","truth",")",".","batch","[","'label","'","]","is","a","dict","where","the","keys","and","values","depend","on","the","losses","used",".","Please","see","labels_losses_and_metrics","and","boxes_losses_and_metrics","member","methods",".","matches",":","Output","of","hungarian","matcher",",","if","available",".","model_params",":","pytree","(","optional",")",";","Parameters","of","the","model",".","Returns",":","total_loss",":","Total","loss","weighted","appropriately","using","self.loss_terms_weights","and","combined","across","all","auxiliary","outputs",".","metrics_dict",":","Individual","loss","terms","with","and","without","weighting","for","logging","purposes","."]
code: def loss_function(  # pytype: disable=signature-mismatch  # overriding-return-type-checks
      self,
      outputs: ArrayDict,
      batch: ArrayDict,
      matches: Optional[jnp.ndarray] = None,
      model_params: Optional[jnp.ndarray] = None,
  ) -> Tuple[jnp.ndarray, MetricsDict]:
      """Loss and metrics function for ObjectDetectionWithMatchingModel.
  
      Args:
        outputs: Model prediction. The exact fields depend on the losses used.
          Please see labels_losses_and_metrics and boxes_losses_and_metrics for
          details.
        batch: Dict that has 'inputs', 'batch_mask' and, 'label' (ground truth).
          batch['label'] is a dict where the keys and values depend on the losses
          used. Please see labels_losses_and_metrics and boxes_losses_and_metrics
          member methods.
        matches: Output of hungarian matcher, if available.
        model_params: pytree (optional); Parameters of the model.
  
      Returns:
        total_loss: Total loss weighted appropriately using
          self.loss_terms_weights and combined across all auxiliary outputs.
        metrics_dict: Individual loss terms with and without weighting for
          logging purposes.
      """
      # Append a padding instance to the inputs. Those are not guaranteed by the
      # input pipeline to always be present.
      batch = batch.copy()
      batch["label"] = batch["label"].copy()
  
      # Append a class label.
      if self.dataset_meta_data["target_is_onehot"]:
          # Shape is [batch, num_instances, num_classes]
          label_shape = batch["label"]["labels"].shape
          num_classes = label_shape[-1]
          instance = jax.nn.one_hot(0, num_classes)
          reshape_shape = (1,) * (len(label_shape) - 1) + (num_classes,)
          broadcast_shape = label_shape[:-2] + (1, num_classes)
          instance = jnp.broadcast_to(
              jnp.reshape(instance, reshape_shape), broadcast_shape
          )
      else:
          instance = jnp.zeros_like(batch["label"]["labels"][..., :1])
      batch["label"]["labels"] = jnp.concatenate(
          [batch["label"]["labels"], instance], axis=-1
      )
  
      # Same for boxes.
      instance = jnp.zeros_like(batch["label"]["boxes"][..., :1, :])
      batch["label"]["boxes"] = jnp.concatenate(
          [batch["label"]["boxes"], instance], axis=-2
      )
  
      if matches is None:
          if "cost" not in outputs:
              cost, n_cols = self.compute_cost_matrix(
                  outputs, batch["label"]
              )  # pytype: disable=wrong-arg-types  # jax-ndarray
          else:
              cost, n_cols = outputs["cost"], outputs.get("cost_n_cols")
          matches = self.matcher(cost, n_cols)
          if "aux_outputs" in outputs:
              matches = [matches]
              for aux_pred in outputs["aux_outputs"]:
                  if "cost" not in outputs:
                      cost, n_cols = self.compute_cost_matrix(
                          aux_pred, batch["label"]
                      )  # pytype: disable=wrong-arg-types  # jax-ndarray
                  else:
                      cost, n_cols = aux_pred["cost"], outputs.get(
                          "cost_n_cols"
                      )
                  matches.append(self.matcher(cost, n_cols))
  
      if not isinstance(matches, (list, tuple)):
          # Ensure matches come as a sequence.
          matches = [matches]
  
      # If the matching is not complete (i.e. the number of tokens is larger than
      # the number of labels) we will pad the matches.
      n = outputs["pred_logits"].shape[-2]
  
      def pad_matches(match):
          b, m = match.shape[0], match.shape[-1]
          if n > m:
  
              def get_unmatched_indices(row, ind):
                  return jax.lax.top_k(
                      jnp.logical_not(row.at[ind].set(1)), k=n - m
                  )
  
              get_unmatched_indices = jax.vmap(get_unmatched_indices)
  
              indices = jnp.zeros((b, n), dtype=jnp.bool_)
              _, indices = get_unmatched_indices(
                  indices, match[:, 0, :]
              )
              indices = jnp.expand_dims(indices, axis=1)
  
              padding = jnp.concatenate(
                  [indices, jnp.full(indices.shape, fill_value=m - 1)],
                  axis=1,
              )
              return jnp.concatenate([match, padding], axis=-1)
          return match
  
      matches = [pad_matches(match) for match in matches]
  
      aux_matches = matches[1:]
      indices = matches[0]
  
      # Computes all the requested losses
      loss_dict = {}
      metrics_dict = {}
      for loss in self.losses_and_metrics:
          loss, metrics = self.get_losses_and_metrics(
              loss, outputs, batch, indices
          )
          loss_dict.update(loss)
          metrics_dict.update(metrics)
  
      # The outputs might have auxiliary predictions from more layers. We process
      # them in the next code block
      if aux_matches:
          for i, aux_outputs in enumerate(outputs["aux_outputs"]):
              indices = aux_matches[i]
              # Computes all the losses for this auxiliary output except class_error
              for loss in self.losses_and_metrics:
                  # Disable class error for loss on labels.
                  kwargs = {"log": False} if loss == "labels" else {}
                  l_dict, m_dict = self.get_losses_and_metrics(
                      loss, aux_outputs, batch, indices, **kwargs
                  )
                  l_dict = {
                      k + f"_aux_{i}": v for k, v in l_dict.items()
                  }
                  loss_dict.update(l_dict)
                  # Add metrics for aux outputs:
                  metrics_dict.update(
                      {k + f"_aux_{i}": v for k, v in m_dict.items()}
                  )
  
      # Compute the total loss by combining loss_dict with loss_terms_weights:
      total_loss = []
      for k, v in loss_dict.items():
          k = k.split("_aux_")[0]  # Remove aux suffix for term weights.
          if k in self.loss_terms_weights:
              total_loss.append(self.loss_terms_weights[k] * v)
      total_loss = sum(total_loss)
  
      if self.config.get("l2_decay_factor") is not None:
          l2_loss = model_utils.l2_regularization(model_params)
          metrics["l2_loss"] = l2_loss
          total_loss = (
              total_loss + 0.5 * self.config.l2_decay_factor * l2_loss
          )
  
      # Process metrics dictionary to generate final unnormalized metrics.
      metrics = self.get_metrics(metrics_dict)
      metrics["total_loss"] = (total_loss, 1)
      return total_loss, 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 9 for calling the BaseModelWithMatching obj's underlying member method code with example usage: obj.loss_function(outputs, batch, matches, model_params) and returns: total_loss,  metrics, match, jnp, jax