class_docstr: Defines commonalities between all semantic segmentation 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 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('simple_cnn_segmentation')
    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","semantic","segmentation","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","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","(","'simple_cnn_segmentation","'",")","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 a callable metric function for global metrics.
  
  The return function implements metrics that require the prediction for the
  entire test/validation dataset in one place and has the following API:
    ```global_metrics_fn(all_confusion_mats, dataset_metadata)```
  If return None, no global metrics will be computed.
docstr_tok: ["Returns","a","callable","metric","function","for","global","metrics",".","The","return","function","implements","metrics","that","require","the","prediction","for","the","entire","test\/validation","dataset","in","one","place","and","has","the","following","API",":","``","`","global_metrics_fn","(","all_confusion_mats",",","dataset_metadata",")","``","`","If","return","None",",","no","global","metrics","will","be","computed","."]
code: def get_global_metrics_fn(self) -> GlobalMetricFn:
      """Returns a callable metric function for global metrics.
  
      The return function implements metrics that require the prediction for the
      entire test/validation dataset in one place and has the following API:
        ```global_metrics_fn(all_confusion_mats, dataset_metadata)```
      If return None, no global metrics will be computed.
      """
      return global_metrics_fn
  
file_path: scenic/scenic/model_lib/base_models/segmentation_model.py
imports: import functools
  from typing import Any, Callable, List, Dict, Optional, Tuple, Union
  from flax.training import common_utils
  from immutabledict import immutabledict
  import jax.numpy
  import 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 4 for calling the SegmentationModel obj's underlying member method code with example usage: obj.get_global_metrics_fn() and returns: global_metrics_fn