class_docstr: Audio Visual Transformer model for multiple n-way classification.
class_docstr_tok: ["Audio","Visual","Transformer","model","for","multiple","n-way","classification","."]
docstr: Return softmax cross entropy loss with an L2 penalty on the weights.
docstr_tok: ["Return","softmax","cross","entropy","loss","with","an","L2","penalty","on","the","weights","."]
code: def loss_function(
      self,
      logits: jnp.ndarray,
      batch: base_model.Batch,
      model_params: Optional[jnp.ndarray] = None,
  ) -> float:
      """Return softmax cross entropy loss with an L2 penalty on the weights."""
      weights = batch.get("batch_mask")
      labels = batch["label"]
  
      assert self.dataset_meta_data.get("target_is_onehot", False)
      if not isinstance(logits, dict):
          all_logits = logits
          logits = {}
          logits["all"] = all_logits
          if isinstance(labels, dict):
              assert "all" in labels, "mixmod must be turned off."
              labels = labels["all"]
          else:
              all_labels = labels
              labels = {}
              labels["all"] = all_labels
  
      sof_ce_loss = []
      for modality in logits:
          if logits[modality].shape[-1] != self.class_splits[-1]:
              raise AssertionError(
                  "Logit dimension must be equal to number of classes"
              )
  
          logit_splits = jnp.split(
              logits[modality], self.class_splits, axis=-1
          )[:-1]
          assert not isinstance(labels[modality], dict), labels.keys()
          labels_splits = jnp.split(
              labels[modality], self.class_splits, axis=-1
          )[:-1]
          label_smoothing = self.config.get("label_smoothing")
  
          sof_ce_losses = [
              base_model_utils.weighted_softmax_cross_entropy(
                  logit_split, labels_split, weights, label_smoothing
              )
              for logit_split, labels_split in zip(
                  logit_splits, labels_splits
              )
          ]
          sof_ce_loss.append(jnp.mean(jnp.array(sof_ce_losses)))
      sof_ce_loss = jnp.mean(jnp.array(sof_ce_loss))
  
      if self.config.get("l2_decay_factor") is None:
          total_loss = sof_ce_loss
      else:
          l2_loss = base_model_utils.l2_regularization(model_params)
          total_loss = (
              sof_ce_loss + 0.5 * self.config.l2_decay_factor * l2_loss
          )
      return total_loss
  
imports: import functools
  from typing import Any, Callable, Dict, Optional, Sequence, Tuple
  from absl import logging
  import flax.linen
  from flax.linen.linear import default_kernel_init
  from immutabledict import immutabledict
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.base_models.classification_model import ClassificationModel
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.baselines import vit
  from scenic.projects.mbt import model_utils
label_desc: Use image node_id 2 for calling the MBTMultiHeadClassificationModel obj's underlying member method code with example usage: obj.loss_function(logits, batch, model_params) and returns: total_loss