class_docstr: Defines metrics/loss among all temporal localization models.
  
  A model is class with three members: get_metrics_fn, loss_fn, & a flax_model.
  
  get_metrics_fn returns a callable function, metric_fn, that calculates the
  metrics and returns a dictionary. The metric function computes f(logits_i,
  batch_i) on a minibatch, it has API:
    ```metric_fn(logits, batch).```
  
  The trainer will then aggregate and compute the mean across all samples
  evaluated.
  
  loss_fn is a function of API
    loss = loss_fn(logits, batch, model_params=None).
  
  This model class defines two losses, sigmoid cross entropy for classification
  and IoU for boundary regression.
class_docstr_tok: ["Defines","metrics\/loss","among","all","temporal","localization","models",".","A","model","is","class","with","three","members",":","get_metrics_fn",",","loss_fn",",","&","a","flax_model",".","get_metrics_fn","returns","a","callable","function",",","metric_fn",",","that","calculates","the","metrics","and","returns","a","dictionary",".","The","metric","function","computes","f","(","logits_i",",","batch_i",")","on","a","minibatch",",","it","has","API",":","``","`","metric_fn","(","logits",",","batch",")",".","``","`","The","trainer","will","then","aggregate","and","compute","the","mean","across","all","samples","evaluated",".","loss_fn","is","a","function","of","API","loss","=","loss_fn","(","logits",",","batch",",","model_params=None",")",".","This","model","class","defines","two","losses",",","sigmoid","cross","entropy","for","classification","and","IoU","for","boundary","regression","."]
docstr: Returns the sum of classification and IoU loss.
  
  Args:
    logits: class logits and predicted start/end time displacements in shape
      (batch_size, num_frames, num_classes * 3) if
      output_per_class_displacements = True, otherwise in shape (batch_size,
      num_frames, num_classes + 2).
    batch: Batch of data.
    model_params: Parameters of the model, for optionally applying
      regularization.
  
  Returns:
    Total loss.
docstr_tok: ["Returns","the","sum","of","classification","and","IoU","loss",".","Args",":","logits",":","class","logits","and","predicted","start\/end","time","displacements","in","shape","(","batch_size",",","num_frames",",","num_classes","*","3",")","if","output_per_class_displacements","=","True",",","otherwise","in","shape","(","batch_size",",","num_frames",",","num_classes","+","2",")",".","batch",":","Batch","of","data",".","model_params",":","Parameters","of","the","model",",","for","optionally","applying","regularization",".","Returns",":","Total","loss","."]
code: def loss_function(
      self,
      logits: jnp.ndarray,
      batch: Batch,
      model_params: Optional[jnp.ndarray] = None,
  ) -> float:
      """Returns the sum of classification and IoU loss.
  
      Args:
        logits: class logits and predicted start/end time displacements in shape
          (batch_size, num_frames, num_classes * 3) if
          output_per_class_displacements = True, otherwise in shape (batch_size,
          num_frames, num_classes + 2).
        batch: Batch of data.
        model_params: Parameters of the model, for optionally applying
          regularization.
  
      Returns:
        Total loss.
      """
      bs, num_frames, _ = logits.shape
      if self.config.get("output_per_class_displacements", True):
          num_classes = logits.shape[-1] // 3
          reshaped_logits = logits.reshape(
              (bs, num_frames, num_classes, 3)
          )
          class_logits = reshaped_logits[..., 0]
          displacements = reshaped_logits[..., 1:]
          gt_displacements = batch["displacements"]
      else:
          num_classes = logits.shape[-1] - 2
          class_logits = logits[..., :num_classes]
          displacements = jnp.expand_dims(
              logits[..., num_classes:], axis=2
          )
          gt_displacements = jnp.expand_dims(
              batch["displacements"], axis=2
          )
      batch_mask = batch["batch_mask"]
      weights = batch_mask[:, None] * batch["inputs"][
          "input_mask"
      ].astype(jnp.float32)
      box_loss_type = self.config.get("box_loss_type", "l1+iou")
      box_loss = weighted_box_regression_loss(
          displacements,
          gt_displacements,
          batch["label"],
          weights=weights,
          loss_type=box_loss_type,
      )
      classification_loss_type = self.config.get(
          "classification_loss_type", "sigmoid"
      )
      if classification_loss_type == "focal":
          classification_loss = weighted_focal_sigmoid_cross_entropy(
              class_logits,
              batch["label"],
              weights=weights,
              label_smoothing=self.config.get("label_smoothing"),
              alpha=self.config.get("focal_loss_alpha", 0.5),
              gamma=self.config.get("focal_loss_gamma", 2.0),
          )
      elif classification_loss_type == "sigmoid":
          classification_loss = (
              model_utils.weighted_sigmoid_cross_entropy(
                  class_logits,
                  batch["label"],
                  weights=weights,
                  label_smoothing=self.config.get("label_smoothing"),
              )
          )
      else:
          raise ValueError(
              f"Unknown loss type: {classification_loss_type}."
          )
      return (
          self.config.get("classification_loss_alpha", 1.0)
          * classification_loss
          + box_loss
      )
  
imports: import functools
  from typing import Any, Callable, Dict, Mapping, Optional, Tuple, Union
  import immutabledict
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  from scenic.projects.unloc import metrics
label_desc: Use image node_id 2 for calling the TemporalLocalizationModel obj's underlying member method code with example usage: obj.loss_function(logits, batch, model_params) and returns: unknown