class_docstr: Sequence to sequence model.
class_docstr_tok: ["Sequence","to","sequence","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].
    targets: Targets to be decoded.
    weights: Indicate which tokens are valid (1) vs padding (0).
    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","]",".","targets",":","Targets","to","be","decoded",".","weights",":","Indicate","which","tokens","are","valid","(","1",")","vs","padding","(","0",")",".","model_params",":","Parameters","of","the","model",",","for","optionally","applying","regularization",".","Returns",":","Total","loss","."]
code: def loss_function(
      self,
      logits: jnp.ndarray,
      targets: jnp.ndarray,
      weights: jnp.ndarray,
      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].
        targets: Targets to be decoded.
        weights: Indicate which tokens are valid (1) vs padding (0).
        model_params: Parameters of the model, for optionally applying
          regularization.
  
      Returns:
        Total loss.
      """
  
      if self.config.get("predict_masked_word", False):
          logits, masked_word_logits = logits
          targets, masked_word_targets = targets
          weights, masked_word_weights = weights
  
      if self.dataset_meta_data.get("target_is_onehot", False):
          one_hot_targets = targets
      else:
          one_hot_targets = common_utils.onehot(
              targets, logits.shape[-1]
          )
  
      sof_ce_loss = model_utils.weighted_softmax_cross_entropy(
          logits,
          one_hot_targets,
          weights,
          label_smoothing=self.config.get("label_smoothing"),
      )
  
      if self.config.get("l2_decay_factor") is None:
          total_loss = sof_ce_loss
      else:
          l2_loss = model_utils.l2_regularization(model_params)
          total_loss = (
              sof_ce_loss + 0.5 * self.config.l2_decay_factor * l2_loss
          )
  
      if self.config.get("predict_masked_word", False):
          mwp_loss = model_utils.weighted_softmax_cross_entropy(
              masked_word_logits,
              common_utils.onehot(
                  masked_word_targets, masked_word_logits.shape[-1]
              ),
              masked_word_weights,
              label_smoothing=self.config.get("label_smoothing"),
          )
          total_loss += mwp_loss * self.config.get(
              "mwp_loss_factor", 1.0
          )
  
      return total_loss
  
imports: from typing import Any, Dict, Optional, Tuple
  from absl import logging
  from flax import linen
  from flax.training import common_utils
  from immutabledict import immutabledict
  import jax
  from jax import lax
  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.projects.mbt import model
  from scenic.projects.mbt.model import temporal_encode
label_desc: Use image node_id 4 for calling the Seq2SeqModel obj's underlying member method code with example usage: obj.loss_function(logits, targets, weights, model_params) and returns: total_loss