class_docstr: Universal Vision Transformer model for multi-label classification task.
docstr: Returns sigmoid cross entropy loss with an L2 penalty on the weights.
  
  Args:
    logits: Output of model in shape [batch, length, num_classes].
    auxiliary_outputs: Output of model auxiliary_outputs, (ponder_times,
      remainders)
    batch: Batch of data that has 'label' and optionally 'batch_mask'.
    model_params: Parameters of the model, for optionally applying
      regularization.
  
  Returns:
    Total loss.
docstr_tok: ["Returns","sigmoid","cross","entropy","loss","with","an","L2","penalty","on","the","weights",".","Args",":","logits",":","Output","of","model","in","shape","[","batch",",","length",",","num_classes","]",".","auxiliary_outputs",":","Output","of","model","auxiliary_outputs",",","(","ponder_times",",","remainders",")","batch",":","Batch","of","data","that","has","'label","'","and","optionally","'batch_mask","'",".","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,
      auxiliary_outputs: Any,
      batch: base_model.Batch,
      model_params: Optional[jnp.ndarray] = None,
  ) -> float:
      """Returns sigmoid cross entropy loss with an L2 penalty on the weights.
  
      Args:
        logits: Output of model in shape [batch, length, num_classes].
        auxiliary_outputs: Output of model auxiliary_outputs, (ponder_times,
          remainders)
        batch: Batch of data that has 'label' and optionally 'batch_mask'.
        model_params: Parameters of the model, for optionally applying
          regularization.
  
      Returns:
        Total loss.
      """
      weights = batch.get("batch_mask")
      ac_config = self.config.model.get("ac_config")
      total_loss = 0.0
  
      if self.dataset_meta_data.get("target_is_onehot", False):
          multihot_target = batch["label"]
      else:
          # This is to support running a multi-label classification model on
          # single-label classification tasks.
          multihot_target = common_utils.onehot(
              batch["label"], logits.shape[-1]
          )
  
      # We calculate the ponder loss only when the ac_config is used.
      if ac_config is not None:
          # Unpack the auxiliary_outputs.
          all_states = auxiliary_outputs[0]
          all_p = auxiliary_outputs[1]
  
          # Calculate different losses for different states
          for i in range(ac_config.act_max_steps):
              sig_ce_loss = model_utils.weighted_sigmoid_cross_entropy(
                  all_states[i],
                  multihot_target,
                  all_p[i]
                  * weights,  # weighted averaging different decisions
                  label_smoothing=self.config.get("label_smoothing"),
              )
              total_loss += sig_ce_loss
          p_g = jnp.zeros(
              [
                  ac_config.act_max_steps,
              ]
          )
          not_halted = 1.0
  
          # Init a geometric prior distribution.
          for i in range(ac_config.act_max_steps):
              # For the last time step, we need to ensure the sum of different
              # steps equal to 1.0.
              if i < ac_config.act_max_steps - 1:
                  p_g = p_g.at[i].set(ac_config.lambda_p * not_halted)
                  not_halted = not_halted * (1 - ac_config.lambda_p)
              else:
                  p_g = p_g.at[-1].set(1.0 - jnp.sum(p_g, axis=0))
          ponder_loss = ponder_loss_fn(all_p, p_g, weights)
          total_loss += ac_config.act_loss_weight * ponder_loss
  
      else:
          # We do not calculate the ponder loss as no act used in config.
          sig_ce_loss = model_utils.weighted_sigmoid_cross_entropy(
              logits,
              multihot_target,
              weights,
              label_smoothing=self.config.get("label_smoothing"),
          )
  
      if self.config.get("l2_decay_factor") is not None:
          l2_loss = model_utils.l2_regularization(model_params)
          total_loss = (
              total_loss + 0.5 * self.config.l2_decay_factor * l2_loss
          )
  
      return total_loss
  
imports: from typing import Any, Optional
  import flax.linen
  from flax.training import common_utils
  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 model_utils
  from scenic.model_lib.base_models.multilabel_classification_model import MultiLabelClassificationModel
  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.baselines.pondernet import layers
label_desc: Use image node_id 2 for calling the PonderViTMultiLabelClassificationModel obj's underlying member method code with example usage: obj.loss_function(logits, auxiliary_outputs, batch, model_params) and returns: total_loss