class_docstr: Abstract class for model with NCR losses.
  
  Supports both softmax-classification and multi-label classification models.
class_docstr_tok: ["Abstract","class","for","model","with","NCR","losses",".","Supports","both","softmax-classification","and","multi-label","classification","models","."]
docstr: Returns softmax cross entropy loss with an L2 penalty 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.
    use_bootstrap: Enable the bootstrap loss term
    loss_weight: Weight for the bootstrap loss term
  
  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","]",".","batch",":","Batch","of","data","that","has","'label","'","and","optionally","'batch_mask","'",".","model_params",":","Parameters","of","the","model",",","for","optionally","applying","regularization",".","use_bootstrap",":","Enable","the","bootstrap","loss","term","loss_weight",":","Weight","for","the","bootstrap","loss","term","Returns",":","Total","loss","."]
code: def ce_loss(
      self,
      logits: Array,
      batch: base_model.Batch,
      model_params: Optional[Array] = None,
      use_bootstrap: bool = False,
      loss_weight: Optional[float] = 1.0,
  ) -> Tuple[float, Dict[str, Array]]:
      """Returns softmax cross entropy loss with an L2 penalty 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.
        use_bootstrap: Enable the bootstrap loss term
        loss_weight: Weight for the bootstrap loss term
  
      Returns:
        Total loss.
      """
      weights = batch.get("batch_mask")
      loss_metrics = {}
  
      if self.dataset_meta_data.get("target_is_onehot", False):
          one_hot_targets = batch["label"]
      else:
          one_hot_targets = common_utils.onehot(
              batch["label"], logits.shape[-1]
          )
  
      softmax_ce_loss = model_utils.weighted_softmax_cross_entropy(
          logits,
          one_hot_targets,
          weights,
          label_smoothing=self.config.get("label_smoothing"),
      )
      loss_metrics["softmax_cross_entropy"] = softmax_ce_loss
  
      if self.config.get("l2_decay_factor") is None:
          total_loss = softmax_ce_loss
      else:
          l2_loss = model_utils.l2_regularization(model_params)
          total_loss = (
              softmax_ce_loss
              + 0.5 * self.config.l2_decay_factor * l2_loss
          )
  
      if use_bootstrap:
          bootstrap_labels = jax.nn.softmax(logits)
          bootstrap_loss = model_utils.weighted_softmax_cross_entropy(
              logits,
              bootstrap_labels,
              weights,
              label_smoothing=self.config.get("label_smoothing"),
          )
          total_loss = (1.0 - loss_weight) * total_loss + (
              loss_weight * bootstrap_loss
          )
  
      # Add the dummy entry for the NCR loss
      loss_metrics["ncr_loss"] = 0.0
      loss_metrics["total_loss"] = total_loss
  
      return total_loss, loss_metrics
  
file_path: scenic/scenic/projects/ncr/base_model.py
imports: import functools
  from typing import Dict, Optional, Tuple, Union
  from flax.training import common_utils
  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 classification_model
  from scenic.model_lib.base_models import model_utils
  from scenic.projects.ncr import loss
label_desc: Use image node_id 2 for calling the NCRModel obj's underlying member method code with example usage: obj.ce_loss(logits, batch, model_params, use_bootstrap, loss_weight) and returns: total_loss,  loss_metrics