code: def weighted_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 weighted macro average of metric.
  
      The weights per class are based on the percentage of positive labels for each
      class.
  
      Args:
        metric_name: Name of metric weighted 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. Note that these weights are
          applied in addition to the weights based on the positive labels for each
          class.
        example_weighted: True if example weights should be applied.
  
      Returns:
        Computation for performing the weighted macro average.
      """
      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,
      )
  
      class_ids = [
          k.class_id for k in sub_keys if k.class_id is not None
      ]
  
      # Compute the weights for labels.
      computations = _class_weights_from_labels(
          class_ids=class_ids,
          eval_config=eval_config,
          model_name=model_name,
          output_name=output_name,
          example_weighted=example_weighted,
      )
      # Class weights metrics are based on a single computation and key.
      class_weights_from_labels_key = computations[0].keys[0]
  
      def result(
          metrics: Dict[metric_types.MetricKey, Any]
      ) -> Dict[metric_types.MetricKey, float]:
          """Returns weighted macro average."""
          class_weights_from_labels = metrics[
              class_weights_from_labels_key
          ]
          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:
                  if (
                      class_weights_from_labels
                      and child_key.sub_key.class_id
                      in class_weights_from_labels
                  ):
                      weight = class_weights_from_labels[offset]
                  if (
                      class_weights
                      and child_key.sub_key.class_id 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}
  
      derived_computation = metric_types.DerivedMetricComputation(
          keys=[key], result=result
      )
      computations.append(derived_computation)
      return computations
  
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 3 for calling a global function with example usage: weighted_macro_average(metric_name, sub_keys, eval_config, model_name, output_name, sub_key, class_weights, example_weighted) and returns: computations, dict