def METHOD_NAME(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()