code: def compute_cost(
      *,
      tgt_labels: jnp.ndarray,
      out_prob: jnp.ndarray,
      tgt_bbox: Optional[jnp.ndarray] = None,
      out_bbox: Optional[jnp.ndarray] = None,
      class_loss_coef: float,
      bbox_loss_coef: Optional[float] = None,
      giou_loss_coef: Optional[float] = None,
      target_is_onehot: bool,
  ) -> jnp.ndarray:
      """Computes cost matrices for a batch of predictions.
  
      Relevant code:
      https://github.com/facebookresearch/detr/blob/647917626d5017e63c1217b99537deb2dcb370d6/models/matcher.py#L35
  
      Args:
        tgt_labels: Class labels of shape [B, M]. If target_is_onehot then it is [B,
          M, C]. Note that the labels corresponding to empty bounding boxes are not
          yet supposed to be filtered out.
        out_prob: Classification probabilities of shape [B, N, C].
        tgt_bbox: Target box coordinates of shape [B, M, 4]. Note that the empty
          bounding boxes are not yet supposed to be filtered out.
        out_bbox: Predicted box coordinates of shape [B, N, 4]
        class_loss_coef: Relative weight of classification loss.
        bbox_loss_coef: Relative weight of bbox loss.
        giou_loss_coef: Relative weight of giou loss.
        target_is_onehot: boolean; Whether targets are one-hot encoded.
  
      Returns:
        A cost matrix [B, N, M].
      """
      if (tgt_bbox is None) != (out_bbox is None):
          raise ValueError(
              "Both `tgt_bbox` and `out_bbox` must be set."
          )
      if (tgt_bbox is not None) and (
          (bbox_loss_coef is None) or (giou_loss_coef is None)
      ):
          raise ValueError(
              "For detection, both `bbox_loss_coef` and `giou_loss_coef`"
              " must be set."
          )
  
      batch_size, max_num_boxes = tgt_labels.shape[:2]
      num_queries = out_prob.shape[1]
      if target_is_onehot:
          mask = tgt_labels[..., 0] == 0  # [B, M]
      else:
          mask = tgt_labels != 0  # [B, M]
  
      # [B, N, M]
      cost_class = -out_prob  # DETR uses -prob for matching.
      max_cost_class = 0.0
  
      # [B, N, M]
      if target_is_onehot:
          cost_class = jnp.einsum(
              "bnl,bml->bnm", cost_class, tgt_labels
          )
      else:
          cost_class = jax.vmap(jnp.take, (0, 0, None))(
              cost_class, tgt_labels, 1
          )
  
      cost = class_loss_coef * cost_class
      cost_upper_bound = max_cost_class
  
      if out_bbox is not None:
          # [B, N, M, 4]
          diff = jnp.abs(out_bbox[:, :, None] - tgt_bbox[:, None, :])
          cost_bbox = jnp.sum(diff, axis=-1)  # [B, N, M]
          cost = cost + bbox_loss_coef * cost_bbox
  
          # Cost_upper_bound is the approximate maximal possible total cost:
          cost_upper_bound = (
              cost_upper_bound + bbox_loss_coef * 4.0
          )  # cost_bbox <= 4
  
          # [B, N, M]
          cost_giou = -box_utils.generalized_box_iou(
              box_utils.box_cxcywh_to_xyxy(out_bbox),
              box_utils.box_cxcywh_to_xyxy(tgt_bbox),
              all_pairs=True,
          )
          cost = cost + giou_loss_coef * cost_giou
  
          # cost_giou < 0, but can be a bit higher in the beginning of training:
          cost_upper_bound = cost_upper_bound + giou_loss_coef * 1.0
  
      # Don't make costs too large w.r.t. the rest to avoid numerical instability.
      mask = mask[:, None]
      cost = cost * mask + (1.0 - mask) * cost_upper_bound
      # Guard against NaNs and Infs.
      cost = jnp.nan_to_num(
          cost,
          nan=cost_upper_bound,
          posinf=cost_upper_bound,
          neginf=cost_upper_bound,
      )
  
      assert cost.shape == (batch_size, num_queries, max_num_boxes)
  
      # Compute the number of unpadded columns for each batch element. It is assumed
      # that all padding is trailing padding.
      n_cols = jnp.where(
          jnp.max(mask, axis=1),
          jnp.expand_dims(jnp.arange(1, max_num_boxes + 1), axis=0),
          0,
      )
      n_cols = jnp.max(n_cols, axis=1)
      return cost, n_cols
  
imports: from typing import Optional
  from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  import jax.numpy
  import numpy
  from scenic.model_lib import matchers
  from scenic.model_lib.base_models import box_utils
  import scipy.optimize
label_desc: Use image node_id 1 for calling a global function with example usage: compute_cost() and returns: cost,  n_cols