def METHOD_NAME(
    metric: Union[
        tf.keras.metrics.Metric, Callable[[], tf.keras.metrics.Metric]
    ]
) -> KerasMetricFinalizer:
  """Creates a finalizer function for the given Keras metric.
  Args:
    metric: An instance of `tf.keras.metrics.Metric` or a no-arg callable that
      constructs a `tf.keras.metrics.Metric`.
  Returns:
    A `tf.function` decorated callable that takes in the unfinalized metric
    values (i.e., tensor values of the variables in `keras_metric.variables`),
    and returns the value of `keras_metric.result()`.
  Raises:
    TypeError: If `metric` is not a `tf.keras.metrics.Metric` and not a no-arg
      callable that returns a `tf.keras.metrics.Metric`.
  """
  @tf.function
  def finalizer(unfinalized_metric_values: list[tf.Tensor]):
    # Construct a new keras metirc here, which is necessary because this
    # `tf.function` may be invoked in a different context as the `model_fn`, and
    # we need the `tf.Variable`s to be created in the current scope in order to
    # use `keras_metric.result()`.
    with tf.init_scope():
      keras_metric = create_keras_metric(metric)
    py_typecheck.check_type(unfinalized_metric_values, list)
    if len(keras_metric.variables) != len(unfinalized_metric_values):
      raise ValueError(
          'The input to the finalizer should be a list of `tf.Tensor`s matching'
          f' the variables of the Keras metric {keras_metric.name}. Expected '
          f'a list of `tf.Tensor`s of length {len(keras_metric.variables)}, '
          f'found a list of length {len(unfinalized_metric_values)}.'
      )
    for v, a in zip(keras_metric.variables, unfinalized_metric_values):
      py_typecheck.check_type(a, tf.Tensor)
      if v.shape != a.shape or v.dtype != a.dtype:
        raise ValueError(
            'The input to the finalizer should be a list of `tf.Tensor`s '
            f'matching the variables of the Keras metric {keras_metric.name}. '
            f'Expected a `tf.Tensor` of shape {v.shape} and dtype {v.dtype!r}, '
            f'found a `tf.Tensor` of shape {a.shape} and dtype {a.dtype!r}.'
        )
      v.assign(a)
    return keras_metric.result()
  return finalizer