def METHOD_NAME(self):
    def metric_fn():
        return MeanMetricAsCompositeTensor()
    # TODO(yuefengz): make _create_per_worker_resources public and get rid
    # of the type_spec hack.
    per_worker_metric = self.cluster_coord._create_per_worker_resources(
        metric_fn
    )
    metric_on_coordinator = metric_fn()
    for metric_remote_value in per_worker_metric._values:
        metric_remote_value._type_spec = metric_on_coordinator._type_spec
    def dataset_fn():
        return tf.data.Dataset.range(1024)
    # TODO(yuefengz): integrate it into model.evaluate.
    @tf.function
    def eval_fn(total_shard, shard_id, metric):
        metric.reset_states()
        dataset_shard = dataset_fn().shard(total_shard, shard_id)
        for i in dataset_shard:
            metric.update_state(i)
        # TODO(yuefengz): we should return the internal state of the metric
        # and then use the combiner API.
        return metric.result()
    total_shards = 128
    result_remote_values = []
    for i in range(total_shards):
        result_remote_values.append(
            self.cluster_coord.schedule(
                eval_fn, args=(total_shards, i, per_worker_metric)
            )
        )
    self._cluster.kill_task("worker", 0)
    self._cluster.kill_task("worker", 1)
    time.sleep(1)
    self._cluster.start_task("worker", 0)
    self._cluster.start_task("worker", 1)
    results = [r.fetch() for r in result_remote_values]
    result = sum(results) / len(results)
    self.assertEqual(result, 511.5)