code: def compute_cost(
      *,
      tgt_labels: jnp.ndarray,
      out_prob: jnp.ndarray,
      tgt_bbox: jnp.ndarray,
      out_bbox: jnp.ndarray,
      alpha: float = 0.25,
      gamma: float = 2.0,
      class_loss_coef: float = 1.0,
      bbox_loss_coef: float = 1.0,
      giou_loss_coef: float = 1.0,
      target_is_onehot: bool = False,
  ) -> Tuple[jnp.ndarray, jnp.ndarray]:
      """Computes cost matrices for DeformableDETR predictions.
  
      Relevant code:
      https://github.com/fundamentalvision/Deformable-DETR/blob/11169a60c33333af00a4849f1808023eba96a931/models/matcher.py#L45
  
      Args:
        tgt_labels: Class labels of shape [bs, ntargets]. If target_is_onehot then
          it is [bs, ntargets, nclasses]. Note that the labels corresponding to
          empty bounding boxes are not yet supposed to be filtered out.
        out_prob: Classification probabilities of shape [bs, nout, nclasses].
        tgt_bbox: Target box coordinates of shape [bs, ntargets, 4]. Note that the
          empty bounding boxes are not yet supposed to be filtered out.
        out_bbox: Predicted box coordinates of shape [bs, nout, 4]
        alpha: Focal loss alpha for class classification loss.
        gamma: Focal loss gamma for class classification loss.
        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: Whether targets are one-hot encoded.
  
      Returns:
        All pairs cost matrix [bs, nout, ntargets].
      """
      # Calculate cost using pred_prob [bs, npreds].
      logfn = lambda x: jnp.log(jnp.clip(x, a_min=1e-8))
      neg_cost_class = (
          (1 - alpha) * (out_prob**gamma) * (-logfn(1 - out_prob))
      )
      pos_cost_class = (
          alpha * ((1 - out_prob) ** gamma) * (-logfn(out_prob))
      )
      cost_class = pos_cost_class - neg_cost_class
  
      # Select class cost for target class [bs, npreds, ntargets].
      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
          )
  
      # Pairwise box l1 [bs, npreds, ntargets, 4].
      diff = jnp.abs(out_bbox[:, :, None] - tgt_bbox[:, None, :])
      # [bs, npreds, ntargets]
      cost_bbox = jnp.sum(diff, axis=-1)
  
      # [bs, npreds, ntargets]
      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,
          eps=1e-8,
      )
  
      total_cost = (
          bbox_loss_coef * cost_bbox
          + class_loss_coef * cost_class
          + giou_loss_coef * cost_giou
      )
  
      # Compute the number of unpadded columns for each batch element. It is assumed
      # that all padding is trailing padding.
      if target_is_onehot:
          tgt_not_padding = tgt_labels[..., 0] == 0
      else:
          tgt_not_padding = tgt_labels != 0
      n_cols = jnp.sum(tgt_not_padding, axis=-1)
      return total_cost, n_cols
  
imports: import math
  from typing import Any, Dict, Optional, Sequence, Tuple, Union
  from absl import logging
  import flax.linen
  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.baselines.deformable_detr.backbone import DeformableDETRBackbone
  from scenic.projects.baselines.deformable_detr.backbone import InputPosEmbeddingSine
  from scenic.projects.baselines.deformable_detr.backbone import mask_for_shape
  from scenic.projects.baselines.deformable_detr.deformable_transformer import BBoxCoordPredictor
  from scenic.projects.baselines.deformable_detr.deformable_transformer import DeformableDETRTransformer
  from scenic.projects.baselines.deformable_detr.deformable_transformer import inverse_sigmoid
  from scenic.projects.baselines.deformable_detr.deformable_transformer import pytorch_kernel_init
label_desc: Use image node_id 1 for calling a global function with example usage: compute_cost() and returns: total_cost,  n_cols