class_docstr: Defines commonalities between all masked self-supervised models.
class_docstr_tok: ["Defines","commonalities","between","all","masked","self-supervised","models","."]
docstr: Returns the (weighted) mean squared error.
  
  Args:
    predictions: Output of model in shape [batch, num_tokens, channels].
    prediction_masks: 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",":","Output","of","model","in","shape","[","batch",",","num_tokens",",","channels","]",".","prediction_masks",":","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  # overriding-parameter-count-checks
      predictions: jnp.ndarray,
      prediction_masks: jnp.ndarray,
      batch: base_model.Batch,
      model_params: Optional[jnp.ndarray] = None,
  ) -> float:
      """Returns the (weighted) mean squared error.
  
      Args:
        predictions: Output of model in shape [batch, num_tokens, channels].
        prediction_masks: 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.
      """
      batch_mask = batch.get("batch_mask")
      if batch_mask is None:
          batch_mask = jnp.ones(prediction_masks.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
  
      feature_target = self.config.masked_feature_loss.target
      if feature_target == FeatureTargets.RGB:
          targets = batch[f"target_{feature_target}"]
      else:
          raise NotImplementedError(
              f"Feature target {feature_target} not implemented."
          )
  
      total_loss = model_utils.weighted_mean_squared_error(
          predictions, 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.shape[-1]
          total_loss = total_loss / output_dimension
  
      if self.config.get("l2_decay_factor"):
          l2_loss = model_utils.l2_regularization(model_params)
          total_loss += 0.5 * self.config.l2_decay_factor * l2_loss
      return total_loss
  
file_path: scenic/scenic/projects/av_mae/base_model.py
imports: import functools
  from typing import Dict, Optional, Tuple, Union
  from immutabledict import immutabledict
  import jax
  import jax.numpy
  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 import regression_model
  from scenic.model_lib.layers import nn_ops
label_desc: Use image node_id 1 for calling the MaskedFeatureRegressionModel obj's underlying member method code with example usage: obj.loss_function(predictions, prediction_masks, batch, model_params) and returns: total_loss