class_docstr: Defines metrics/loss among all moment retrieval 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","moment","retrieval","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 a callable metric function for the model.
  
  Args:
    split: The split for which we calculate the metrics. It should be one of
      the ['train',  'validation', 'test'].
  Returns: A metric function with the following API: ```metrics_fn(logits,
    batch)```
docstr_tok: ["Returns","a","callable","metric","function","for","the","model",".","Args",":","split",":","The","split","for","which","we","calculate","the","metrics",".","It","should","be","one","of","the","[","'train","'",",","'validation","'",",","'test","'","]",".","Returns",":","A","metric","function","with","the","following","API",":","``","`","metrics_fn","(","logits",",","batch",")","``","`"]
code: def get_metrics_fn(self, split: Optional[str] = None):
      """Returns a callable metric function for the model.
  
      Args:
        split: The split for which we calculate the metrics. It should be one of
          the ['train',  'validation', 'test'].
      Returns: A metric function with the following API: ```metrics_fn(logits,
        batch)```
      """
      del split  # For all splits, we return the same metric functions.
      cls_loss_type = self.config.get(
          "classification_loss_type", "sigmoid"
      )
      box_loss_type = self.config.get("box_loss_type", "l1+iou")
      box_loss_types = box_loss_type.split("+")
      cls_metrics = (
          _MOMENT_RETRIEVAL_FOCAL_LOSS_CLASSIFICATION_METRICS
          if cls_loss_type == "focal"
          else _MOMENT_RETRIEVAL_SIGMOID_LOSS_CLASSIFICATION_METRICS
      )
      box_regression_metrics = dict(
          _MOMENT_RETRIEVAL_BOX_REGRESSION_METRICS
      )
      for weight_and_type in box_loss_types:
          loss_type = weight_and_type.split("*")[-1]
          box_regression_metrics[f"{loss_type}_loss"] = (
              functools.partial(
                  temporal_localization_base_model.weighted_unnormalized_box_regression_loss,
                  loss_type=loss_type,
              ),
              temporal_localization_base_model.num_positive_frames,
          )
      return functools.partial(
          moment_retrieval_metrics_function,
          config=self.config,
          classification_metrics=cls_metrics,
          box_regression_metrics=box_regression_metrics,
      )
  
imports: import functools
  from typing import Any, Dict, Mapping, Optional, Tuple, Union
  import immutabledict
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  from scenic.projects.unloc import temporal_localization_base_model
label_desc: Use image node_id 1 for calling the MomentRetrievalModel obj's underlying member method code with example usage: obj.get_metrics_fn(split) and returns: functools