code: def macro_average(
      metric_name: str,
      sub_keys: Iterable[metric_types.SubKey],
      eval_config: Optional[config_pb2.EvalConfig] = None,
      model_name: str = "",
      output_name: str = "",
      sub_key: Optional[metric_types.SubKey] = None,
      class_weights: Optional[Dict[int, float]] = None,
      example_weighted: bool = False,
  ) -> metric_types.MetricComputations:
      """Returns metric computations for computing macro average of given metric.
  
      Args:
        metric_name: Name of underlying metric average is being computed for.
        sub_keys: Sub keys used to compute the metric (e.g. class_ids, etc).
        eval_config: Eval config.
        model_name: Optional model name.
        output_name: Optional output name.
        sub_key: Optional sub key associated with aggregation metric (e.g. top_k).
        class_weights: Optional class weights to apply. Required if sub_key is not
          provided. If class_weights are provided, but a sub_key.class_id (if
          sub_key is None) or sub_key.k (if sub_key is top_k) is not set or not
          found in the dictionary then 0.0 is assumed.
        example_weighted: True if example weights should be applied.
  
      Returns:
        Computation for performing the macro average.
      """
      del eval_config
  
      key = metric_types.MetricKey(
          name=metric_name,
          model_name=model_name,
          output_name=output_name,
          sub_key=sub_key,
          aggregation_type=metric_types.AggregationType(
              macro_average=True
          ),
          example_weighted=example_weighted,
      )
  
      def result(
          metrics: Dict[metric_types.MetricKey, float]
      ) -> Dict[metric_types.MetricKey, float]:
          """Returns macro average."""
          total_value = 0.0
          total_weight = 0.0
          for sub_key in sub_keys:
              child_key = metric_types.MetricKey(
                  name=metric_name,
                  model_name=model_name,
                  output_name=output_name,
                  sub_key=sub_key,
                  example_weighted=example_weighted,
              )
              if child_key not in metrics:
                  # Use private name if not found under metric name
                  child_key = metric_types.MetricKey(
                      name="_" + metric_name,
                      model_name=model_name,
                      output_name=output_name,
                      sub_key=sub_key,
                      example_weighted=example_weighted,
                  )
              weight = 1.0 if not class_weights else 0.0
              offset = None
              if (
                  child_key.sub_key is not None
                  and child_key.sub_key.class_id is not None
              ):
                  offset = child_key.sub_key.class_id
              elif (
                  child_key.sub_key is not None
                  and child_key.sub_key.k is not None
              ):
                  offset = child_key.sub_key.k
              if offset is not None and offset in class_weights:
                  weight = class_weights[offset]
              total_value += _to_float(metrics[child_key]) * weight
              total_weight += weight
          average = (
              total_value / total_weight
              if total_weight
              else float("nan")
          )
          return {key: average}
  
      return [
          metric_types.DerivedMetricComputation(
              keys=[key], result=result
          )
      ]
  
file_path: model-analysis/tensorflow_model_analysis/metrics/aggregation.py
imports: from typing import Any, Dict, Iterable, List, Optional
  import apache_beam
  from tensorflow_model_analysis.metrics import metric_types
  from tensorflow_model_analysis.metrics import metric_util
  from tensorflow_model_analysis.proto import config_pb2
label_desc: Use image node_id 2 for calling a global function with example usage: macro_average(metric_name, sub_keys, eval_config, model_name, output_name, sub_key, class_weights, example_weighted) and returns: list, dict