class_docstr: Defines commonalities between all models.
  
  A model is class with three members: get_metrics_fn, loss_fn, and 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 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","models",".","A","model","is","class","with","three","members",":","get_metrics_fn",",","loss_fn",",","and","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","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",",","...",")","``","`"]
code: def __init__(
      self,
      config: Optional[ml_collections.ConfigDict],
      dataset_meta_data: Dict[str, Any],
  ) -> None:
      if config is None:
          logging.warning(
              "You are creating the model with default config."
          )
          config = self.default_flax_model_config()
      self.config = config
      self.dataset_meta_data = dataset_meta_data
      self.flax_model = self.build_flax_model()
  
file_path: scenic/scenic/model_lib/base_models/base_model.py
imports: from typing import Any, Callable, Dict, Mapping, Optional, Tuple
  from absl import logging
  import flax.linen
  from flax.training import common_utils
  import jax.numpy
  import ml_collections
label_desc: Use image node_id 1 to create a new BaseModel object  with example: obj = BaseModel(config, dataset_meta_data)