class_docstr: A object which creates evaluation metrics for Subnetworks.
class_docstr_tok: ["A","object","which","creates","evaluation","metrics","for","Subnetworks","."]
docstr: Wraps the given metric_fn with a template so it's Variables are shared.
  
  Hooks on TPU cannot depend on any graph Tensors. Instead the eval metrics
  returned by metric_fn are stored in Variables. These variables are later
  read from the evaluation hooks which run on the host CPU.
  
  Args:
    metric_fn: The function to wrap with a template.
  
  Returns:
    The original metric_fn wrapped with a template function.
docstr_tok: ["Wraps","the","given","metric_fn","with","a","template","so","it","'s","Variables","are","shared",".","Hooks","on","TPU","can","not","depend","on","any","graph","Tensors",".","Instead","the","eval","metrics","returned","by","metric_fn","are","stored","in","Variables",".","These","variables","are","later","read","from","the","evaluation","hooks","which","run","on","the","host","CPU",".","Args",":","metric_fn",":","The","function","to","wrap","with","a","template",".","Returns",":","The","original","metric_fn","wrapped","with","a","template","function","."]
code: def _templatize_metric_fn(self, metric_fn):
      """Wraps the given metric_fn with a template so it's Variables are shared.
  
      Hooks on TPU cannot depend on any graph Tensors. Instead the eval metrics
      returned by metric_fn are stored in Variables. These variables are later
      read from the evaluation hooks which run on the host CPU.
  
      Args:
        metric_fn: The function to wrap with a template.
  
      Returns:
        The original metric_fn wrapped with a template function.
      """
  
      def _metric_fn(*args, **kwargs):
          """The wrapping function to be returned."""
  
          # We can only be passed in either a dict or a list of tensors.
          args = args if args else kwargs
          metrics = _call_eval_metrics((metric_fn, args))
          if not self._use_tpu:
              return metrics
  
          logging.log_first_n(
              logging.INFO,
              "Writing eval metrics to variables for TPU",
              1,
          )
          wrapped_metrics = {}
          for i, key in enumerate(sorted(metrics)):
              tensor, op = tf_compat.metric_op(metrics[key])
              # key cannot be in var name since it may contain illegal chars.
              var = tf_compat.v1.get_variable(
                  "metric_{}".format(i),
                  shape=tensor.shape,
                  dtype=tensor.dtype,
                  trainable=False,
                  initializer=tf_compat.v1.zeros_initializer(),
                  collections=[tf_compat.v1.GraphKeys.LOCAL_VARIABLES],
              )
              if (
                  isinstance(op, tf.Operation)
                  or op.shape != tensor.shape
              ):
                  with tf.control_dependencies([op]):
                      op = var.assign(tensor)
              metric = (var, var.assign(op))
              wrapped_metrics[key] = metric
          return wrapped_metrics
  
      return tf_compat.v1.make_template(
          "metric_fn_template", _metric_fn
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import inspect
  from absl import logging
  from adanet import tf_compat
  import six
  import tensorflow.compat.v2
label_desc: Use image node_id 3 for calling the _SubnetworkMetrics obj's underlying member method code with example usage: obj._templatize_metric_fn(metric_fn) and returns: tf_compat, wrapped_metrics, metrics