class_docstr: PolyVit base model.
docstr: Returns softmax cross entropy loss with an L2 penalty on the weights.
  
  Args:
    logits: Output of model in shape [batch, length, num_classes].
    batch: Batch of data that has 'label' and optionally 'batch_mask'.
    dataset_name: The name of the dataset used for the task.
    model_params: Parameters of the model, for optionally applying
      regularization.
  
  Returns:
    Total loss.
docstr_tok: ["Returns","softmax","cross","entropy","loss","with","an","L2","penalty","on","the","weights",".","Args",":","logits",":","Output","of","model","in","shape","[","batch",",","length",",","num_classes","]",".","batch",":","Batch","of","data","that","has","'label","'","and","optionally","'batch_mask","'",".","dataset_name",":","The","name","of","the","dataset","used","for","the","task",".","model_params",":","Parameters","of","the","model",",","for","optionally","applying","regularization",".","Returns",":","Total","loss","."]
code: def loss_function(
      self,  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
      logits: jnp.ndarray,
      batch: base_model.Batch,
      dataset_name: str,
      model_params: Optional[jnp.ndarray] = None,
  ) -> float:
      """Returns softmax cross entropy loss with an L2 penalty on the weights.
  
      Args:
        logits: Output of model in shape [batch, length, num_classes].
        batch: Batch of data that has 'label' and optionally 'batch_mask'.
        dataset_name: The name of the dataset used for the task.
        model_params: Parameters of the model, for optionally applying
          regularization.
  
      Returns:
        Total loss.
      """
      weights = batch.get("batch_mask")
      class_splits, _ = self._get_splits()
  
      current_dataset_meta_data = self.dataset_meta_data[dataset_name]
      task = current_dataset_meta_data["task"]
  
      if current_dataset_meta_data.get("target_is_onehot", False):
          one_or_multi_hot_targets = batch["label"]
      elif task == Task.MULTIHEADLABEL:
          raise ValueError("Target labels should be one-hot.")
      else:
          one_or_multi_hot_targets = common_utils.onehot(
              batch["label"], logits.shape[-1]
          )
  
      if self.config.get("label_smoothing_params") is not None:
          label_smoothing = self.config.label_smoothing_params.get(
              dataset_name, None
          )
      else:
          label_smoothing = self.config.get("label_smoothing")
  
      kwargs = {"label_smoothing": label_smoothing}
  
      if task == Task.MULTIHEADLABEL:
          kwargs["class_splits"] = class_splits[dataset_name]
  
      if task == Task.BOW:
          # one_or_multi_hot_targets is initially one-hot of shape (bs, len,
          # vocab_size), while we actually need multi-hot of shape (bs, vocab_size).
          one_or_multi_hot_targets = one_or_multi_hot_targets.max(
              axis=-2
          )
  
      loss = _LOSS_FUNCTIONS[task](
          logits, one_or_multi_hot_targets, weights, **kwargs
      )
  
      config_loss_weight = self.config.get(
          "loss_weights", ml_collections.ConfigDict()
      )
      loss_weight = config_loss_weight.get(dataset_name, 1.0)
  
      loss = loss * loss_weight
  
      if self.config.get("l2_decay_factor") is None:
          total_loss = loss
      else:
          l2_loss = model_utils.l2_regularization(model_params)
          total_loss = (
              loss + 0.5 * self.config.l2_decay_factor * l2_loss
          )
      return total_loss
  
imports: import functools
  from typing import Any, Dict, Optional, Tuple, List
  from absl import logging
  from flax.training import common_utils
  from immutabledict 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.model_lib.base_models.classification_model import classification_metrics_function
  from scenic.model_lib.base_models.encoder_decoder_model import encoder_decoder_metrics_function
  from scenic.model_lib.base_models.multilabel_classification_model import multilabel_classification_metrics_function
  from scenic.projects.vivit import model_utils
label_desc: Use image node_id 4 for calling the PolyVitBaseModel obj's underlying member method code with example usage: obj.loss_function(logits, batch, dataset_name, model_params) and returns: total_loss