code: def loss_labels(
      *,
      pred_logits: jnp.ndarray,
      tgt_labels: jnp.ndarray,
      indices: jnp.ndarray,
      alpha: float = 0.25,
      gamma: float = 2.0,
      class_loss_coef: float = 1.0,
      target_is_onehot: bool = False,
  ) -> ArrayDict:
      """Calculate DeformableDETR classification loss.
  
      Args:
        pred_logits: [bs, n_preds, n_classes].
        tgt_labels: [bs, n_max_targets].
        indices: [bs, 2, min(n_preds, n_max_targets)].
        alpha: Focal loss alpha.
        gamma: Focal loss gamma.
        class_loss_coef: Classification loss coefficient.
        target_is_onehot: Tgt is [bs, n_max_targets, n_classes]
  
      Returns:
        `loss_class`: Classification loss with coefficient applied.
      """
      # Apply the permutation communicated by indices.
      pred_logits = model_utils.simple_gather(
          pred_logits, indices[..., 0, :]
      )
      tgt_labels = model_utils.simple_gather(
          tgt_labels, indices[..., 1, :]
      )
  
      if target_is_onehot:
          tgt_labels_onehot = tgt_labels
      else:
          nclasses = pred_logits.shape[-1]
          tgt_labels_onehot = jnp.where(
              tgt_labels == 0, nclasses, tgt_labels
          )
          tgt_labels_onehot = jax.nn.one_hot(
              tgt_labels_onehot, nclasses
          )
  
      loss = model_utils.focal_sigmoid_cross_entropy(
          pred_logits, tgt_labels_onehot, alpha=alpha, gamma=gamma
      )
      loss = loss.mean(1).sum() * pred_logits.shape[1]
      loss = class_loss_coef * loss
      return {"loss_class": loss}
  
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 2 for calling a global function with example usage: loss_labels() and returns: dict