class_docstr: Defines commonalities between all classification models.
  
  A model is class with three members: get_metrics_fn, loss_fn, & a flax_model.
  
  get_metrics_fn returns a callable function, metric_fn, that calculates the
  metrics and returns a dictionary. The metric function computes f(x_i, y_i) on
  a minibatch, it has API:
    ```metric_fn(logits, label, weights).```
  
  The trainer will then aggregate and compute the mean across all samples
  evaluated.
  
  loss_fn is a function of API
    loss = loss_fn(logits, batch, model_params=None).
  
  This model class defines a softmax_cross_entropy_loss with weight decay,
  where the weight decay factor is determined by config.l2_decay_factor.
  
  flax_model is returned from the build_flax_model function. A typical
  usage pattern will be:
    ```
    model_cls = model_lib.models.get_model_cls('fully_connected_classification')
    model = model_cls(config, dataset.meta_data)
    flax_model = model.build_flax_model
    dummy_input = jnp.zeros(input_shape, model_input_dtype)
    model_state, params = flax_model.init(
        rng, dummy_input, train=False).pop('params')
    ```
  And this is how to call the model:
    variables = {'params': params, **model_state}
    logits, new_model_state = flax_model.apply(variables, inputs, ...)
    ```
class_docstr_tok: ["Defines","commonalities","between","all","classification","models",".","A","model","is","class","with","three","members",":","get_metrics_fn",",","loss_fn",",","&","a","flax_model",".","get_metrics_fn","returns","a","callable","function",",","metric_fn",",","that","calculates","the","metrics","and","returns","a","dictionary",".","The","metric","function","computes","f","(","x_i",",","y_i",")","on","a","minibatch",",","it","has","API",":","``","`","metric_fn","(","logits",",","label",",","weights",")",".","``","`","The","trainer","will","then","aggregate","and","compute","the","mean","across","all","samples","evaluated",".","loss_fn","is","a","function","of","API","loss","=","loss_fn","(","logits",",","batch",",","model_params=None",")",".","This","model","class","defines","a","softmax_cross_entropy_loss","with","weight","decay",",","where","the","weight","decay","factor","is","determined","by","config.l2_decay_factor",".","flax_model","is","returned","from","the","build_flax_model","function",".","A","typical","usage","pattern","will","be",":","``","`","model_cls","=","model_lib.models.get_model_cls","(","'fully_connected_classification","'",")","model","=","model_cls","(","config",",","dataset.meta_data",")","flax_model","=","model.build_flax_model","dummy_input","=","jnp.zeros","(","input_shape",",","model_input_dtype",")","model_state",",","params","=","flax_model.init","(","rng",",","dummy_input",",","train=False",")",".pop","(","'params","'",")","``","`","And","this","is","how","to","call","the","model",":","variables","=","{","'params","'",":","params",",","*","*","model_state","}","logits",",","new_model_state","=","flax_model.apply","(","variables",",","inputs",",","...",")","``","`"]
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.
  
  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",".","Returns",":","Total","loss","."]
code: def loss_function(
      self,
      logits: jnp.ndarray,
      batch: base_model.Batch,
      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].
        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")
  
      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]
          )
  
      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
          )
      return total_loss
  
file_path: scenic/scenic/model_lib/base_models/classification_model.py
imports: import functools
  from typing import Dict, Optional, Tuple, Union
  from flax.training import common_utils
  from immutabledict import immutabledict
  import jax.numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
label_desc: Use image node_id 3 for calling the ClassificationModel obj's underlying member method code with example usage: obj.loss_function(logits, batch, model_params) and returns: total_loss