class_docstr: Base model for object detection with matching.
class_docstr_tok: ["Base","model","for","object","detection","with","matching","."]
docstr: Loss and metrics function for matching models.
  
  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 a matcher [B, 2, M]. If not provided, will be computed.
    model_params: pytree (optional); Parameters of the model.
  
  Returns:
    total_loss: Total loss weighted appropriately using
      self.loss_terms_weights.
    metrics_dict: Individual loss terms with and without weighting for
      logging purposes.
docstr_tok: ["Loss","and","metrics","function","for","matching","models",".","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","a","matcher","[","B",",","2",",","M","]",".","If","not","provided",",","will","be","computed",".","model_params",":","pytree","(","optional",")",";","Parameters","of","the","model",".","Returns",":","total_loss",":","Total","loss","weighted","appropriately","using","self.loss_terms_weights",".","metrics_dict",":","Individual","loss","terms","with","and","without","weighting","for","logging","purposes","."]
code: def loss_function(  # pytype: disable=signature-mismatch
      self,
      outputs: ArrayDict,
      batch: ArrayDict,
      matches: Optional[jnp.ndarray] = None,
      model_params: Optional[PyTree] = None,
  ) -> Tuple[jnp.ndarray, MetricsDict]:
      """Loss and metrics function for matching models.
  
      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 a matcher [B, 2, M]. If not provided, will be computed.
        model_params: pytree (optional); Parameters of the model.
  
      Returns:
        total_loss: Total loss weighted appropriately using
          self.loss_terms_weights.
        metrics_dict: Individual loss terms with and without weighting for
          logging purposes.
      """
      batch = batch.copy()
      batch["label"] = batch["label"].copy()
  
      # Append an instance with "padding" label (i.e., "0" as the class label).
      # Shape is [batch, num_instances, num_classes]. This is necessary because
      # the matching code requires at least one padding instance, to which
      # unmatched instances will be assigned.
      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
      )
      batch["label"]["labels"] = jnp.concatenate(
          [batch["label"]["labels"], instance], axis=-2
      )
      if "boxes" in batch["label"]:
          instance = jnp.zeros_like(batch["label"]["boxes"][..., :1, :])
          batch["label"]["boxes"] = jnp.concatenate(
              [batch["label"]["boxes"], instance], axis=-2
          )
  
      # Compute matches if not provided.
      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 not isinstance(matches, (list, tuple)):
          # Ensure matches come as a sequence.
          matches = [matches]
  
      # Pad matches if the matching is not complete (i.e. the number of
      # predicted instances is larger than the number of gt instances).
      num_pred = outputs["pred_logits"].shape[-2]
  
      def pad_matches(match):
          batch_size, _, num_matched = match.shape  # [B, 2, M]
          if num_pred > num_matched:
  
              def get_unmatched_indices(row, ind):
                  return jax.lax.top_k(
                      jnp.logical_not(row.at[ind].set(1)),
                      k=num_pred - num_matched,
                  )
  
              get_unmatched_indices = jax.vmap(get_unmatched_indices)
  
              indices = jnp.zeros(
                  (batch_size, num_pred), 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=num_matched - 1
                      ),
                  ],
                  axis=1,
              )
              return jnp.concatenate([match, padding], axis=-1)
          return match
  
      matches = [pad_matches(match) for match in matches]
  
      indices = matches[0]
  
      # Compute all the requested losses and metrics.
      loss_dict = {}
      metrics_dict = {}
      for loss_name in self.losses_and_metrics:
          loss, metrics = self.get_losses_and_metrics(
              loss_name, outputs, batch, indices
          )
          loss_dict.update(loss)
          metrics_dict.update(metrics)
  
      # Compute the total loss by combining loss_dict with loss_terms_weights.
      total_loss = []
      for k, v in loss_dict.items():
          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_dict["l2_loss"] = (l2_loss, 1)
          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
  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 8 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