class_docstr: AdaTape Transformer model for multi-label classification task.
class_docstr_tok: ["AdaTape","Transformer","model","for","multi-label","classification","task","."]
docstr: Returns sigmoid cross entropy loss with an L2 penalty (and ponder loss) on the weights.
  
  Args:
    logits: Output of model in shape [batch, length, num_classes].
    batch: Batch of data that has 'label' and optionally 'batch_mask'.
    model_params: Parameters of the model, for optionally applying
      regularization.
    auxiliary_outputs: Output of model auxiliary_outputs.
  
  Returns:
    Total loss.
docstr_tok: ["Returns","sigmoid","cross","entropy","loss","with","an","L2","penalty","(","and","ponder","loss",")","on","the","weights",".","Args",":","logits",":","Output","of","model","in","shape","[","batch",",","length",",","num_classes","]",".","batch",":","Batch","of","data","that","has","'label","'","and","optionally","'batch_mask","'",".","model_params",":","Parameters","of","the","model",",","for","optionally","applying","regularization",".","auxiliary_outputs",":","Output","of","model","auxiliary_outputs",".","Returns",":","Total","loss","."]
code: def loss_function(
      self,
      logits: jnp.ndarray,
      batch: base_model.Batch,
      model_params: Optional[jnp.ndarray] = None,
      auxiliary_outputs: Any = None,
  ) -> float:
      """Returns sigmoid cross entropy loss with an L2 penalty (and ponder loss) on the weights.
  
      Args:
        logits: Output of model in shape [batch, length, num_classes].
        batch: Batch of data that has 'label' and optionally 'batch_mask'.
        model_params: Parameters of the model, for optionally applying
          regularization.
        auxiliary_outputs: Output of model auxiliary_outputs.
  
      Returns:
        Total loss.
      """
      weights = batch.get("batch_mask")
  
      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]
          )
  
      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 None:
          total_loss = sig_ce_loss
      else:
          l2_loss = model_utils.l2_regularization(model_params)
          total_loss = (
              sig_ce_loss + 0.5 * self.config.l2_decay_factor * l2_loss
          )
      act_config = self.config.model.ac_config.get(
          "dynamic_tape_length"
      )
      if (act_config is not None) and (
          act_config.act_loss_weight > 0.0
      ):
          ponder_loss = ponder_loss_fn(auxiliary_outputs[1], weights)
          total_loss += act_config.act_loss_weight * ponder_loss
      return total_loss
  
imports: from typing import Any, Optional
  from absl import logging
  import flax
  import flax.linen
  from flax.training import common_utils
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.common_lib import debug_utils
  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 nn_layers
  from scenic.projects.adatape import layers
  from scenic.projects.baselines import vit
label_desc: Use image node_id 2 for calling the AdaTapeMultiLabelClassificationModel obj's underlying member method code with example usage: obj.loss_function(logits, batch, model_params, auxiliary_outputs) and returns: total_loss