class_docstr: ViViT model for multi-modalitymasked autoencoder pretraining.
docstr: Returns the (weighted) mean squared error.
  
  Args:
    predictions: Dictionary with the output of model in shape
      [batch, num_tokens, channels].
    prediction_masks: Dictionary with the tokens to compute the loss on.
      Shape is [batch, num_tokens]
    batch: Batch (dict) with keys 'targets' and optionally 'batch_mask'.
    model_params: Parameters of the model, for optionally applying
      L2 regularization.
  
  Returns:
    The (weighted) mean squared error.
docstr_tok: ["Returns","the","(","weighted",")","mean","squared","error",".","Args",":","predictions",":","Dictionary","with","the","output","of","model","in","shape","[","batch",",","num_tokens",",","channels","]",".","prediction_masks",":","Dictionary","with","the","tokens","to","compute","the","loss","on",".","Shape","is","[","batch",",","num_tokens","]","batch",":","Batch","(","dict",")","with","keys","'targets","'","and","optionally","'batch_mask","'",".","model_params",":","Parameters","of","the","model",",","for","optionally","applying","L2","regularization",".","Returns",":","The","(","weighted",")","mean","squared","error","."]
code: def loss_function(
      self,  # pytype: disable=signature-mismatch  # jax-ndarray
      predictions: ArrayDict,
      prediction_masks: ArrayDict,
      batch: base_model_lib.Batch,
      model_params: Optional[jnp.ndarray] = None,
  ) -> float:
      """Returns the (weighted) mean squared error.
  
      Args:
        predictions: Dictionary with the output of model in shape
          [batch, num_tokens, channels].
        prediction_masks: Dictionary with the tokens to compute the loss on.
          Shape is [batch, num_tokens]
        batch: Batch (dict) with keys 'targets' and optionally 'batch_mask'.
        model_params: Parameters of the model, for optionally applying
          L2 regularization.
  
      Returns:
        The (weighted) mean squared error.
      """
  
      def get_loss_weights(target_name_):
          batch_mask = batch.get("batch_mask")
          if batch_mask is None:
              batch_mask = jnp.ones(
                  prediction_masks[target_name_].shape
              )
          if batch_mask.ndim == 1:
              batch_mask = jnp.expand_dims(batch_mask, axis=-1)
          if self.config.masked_feature_loss.get(
              "loss_unmasked_tokens", False
          ):
              loss_weights = batch_mask
          else:
              loss_weights = batch_mask * prediction_masks[target_name_]
  
          return loss_weights
  
      total_loss = 0.0
      for target_name in self.config.masked_feature_loss.target:
          targets = batch["targets"][target_name]
          loss_weights = get_loss_weights(target_name)
          loss = model_utils_lib.weighted_mean_squared_error(
              predictions[target_name], targets, loss_weights, axis=-1
          )
  
          # Mean squared error is normalised by the number of tokens.
          # If this option is enabled, we normalise further by the number
          # of features we are regressing to.
          if self.config.masked_feature_loss.get(
              "normalise_by_output_dimension", False
          ):
              output_dimension = predictions[target_name].shape[-1]
              loss = loss / output_dimension
          total_loss = total_loss + (
              loss
              * self.config.masked_feature_loss.modality_weight[
                  target_name
              ]
          )
  
      if self.config.get("l2_decay_factor"):
          l2_loss = model_utils_lib.l2_regularization(model_params)
          total_loss += 0.5 * self.config.l2_decay_factor * l2_loss
      return total_loss
  
imports: import functools
  from typing import Any, Dict, Optional, Tuple
  import flax.linen
  import jax
  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.model_lib.layers import nn_layers
  from scenic.projects.av_mae import base_model
  from scenic.projects.av_mae import model_utils
  from scenic.projects.mbt import model
  from scenic.projects.vivit import model
  from scenic.projects.vivit import model_utils
label_desc: Use image node_id 3 for calling the ViViTMultiMaskedAutoencoderModel obj's underlying member method code with example usage: obj.loss_function(predictions, prediction_masks, batch, model_params) and returns: total_loss, loss_weights