def METHOD_NAME(argv):
  """Run a CNN model on MNIST data to demonstrate DistributedStrategies."""
  del argv  # Unused.
  tf.disable_v2_behavior()
  cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
      tpu=FLAGS.tpu)
  strategy = tf.contrib.distribute.TPUStrategy(cluster_resolver)
  # TODO(rxsang): This doesn't work on cloud for some reason
  strategy.extended.experimental_enable_get_next_as_optional = False
  with strategy.scope():
    train_ds, test_ds = mnist_datasets()
    train_ds = train_ds.shuffle(NUM_TRAIN_IMAGES).batch(FLAGS.batch_size)
    test_ds = test_ds.batch(FLAGS.batch_size)
    model = create_model()
    optimizer = tf.keras.optimizers.SGD(FLAGS.learning_rate, FLAGS.momentum)
    training_loss = tf.keras.metrics.Mean("training_loss", dtype=tf.float32)
    training_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
        "training_accuracy", dtype=tf.float32)
    test_loss = tf.keras.metrics.Mean("test_loss", dtype=tf.float32)
    test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
        "test_accuracy", dtype=tf.float32)
    def train_step(inputs):  # pylint: disable=missing-docstring
      images, labels = inputs
      with tf.GradientTape() as tape:
        logits = model(images, training=True)
        loss = compute_loss(logits, labels)
      grads = tape.gradient(loss, model.variables)
      update_vars = optimizer.apply_gradients(zip(grads, model.variables))
      update_loss = training_loss.update_state(loss)
      update_accuracy = training_accuracy.update_state(labels, logits)
      with tf.control_dependencies([update_vars, update_loss, update_accuracy]):
        return tf.identity(loss)
    def test_step(inputs):
      images, labels = inputs
      logits = model(images, training=False)
      loss = compute_loss(logits, labels)
      update_loss = test_loss.update_state(loss)
      update_accuracy = test_accuracy.update_state(labels, logits)
      with tf.control_dependencies([update_loss, update_accuracy]):
        return tf.identity(loss)
    train_iterator = strategy.experimental_distribute_dataset(
        train_ds).make_initializable_iterator()
    test_iterator = strategy.experimental_distribute_dataset(
        test_ds).make_initializable_iterator()
    dist_train = strategy.experimental_local_results(
        strategy.experimental_run_v2(train_step, args=(next(train_iterator),)))
    dist_test = strategy.experimental_local_results(
        strategy.experimental_run_v2(test_step, args=(next(test_iterator),)))
    training_loss_result = training_loss.result()
    training_accuracy_result = training_accuracy.result()
    test_loss_result = test_loss.result()
    test_accuracy_result = test_accuracy.result()
    tf.contrib.distribute.initialize_tpu_system(cluster_resolver)
    train_iterator_init = train_iterator.initialize()
    test_iterator_init = test_iterator.initialize()
    all_variables = (
        tf.global_variables() +
        training_loss.variables +
        training_accuracy.variables +
        test_loss.variables +
        test_accuracy.variables)
    config = tf.ConfigProto()
    cluster_spec = cluster_resolver.cluster_spec()
    if cluster_spec:
      config.cluster_def.CopyFrom(cluster_spec.as_cluster_def())
    with tf.Session(cluster_resolver.master(), config=config) as session:
      session.run([v.initializer for v in all_variables])
      for epoch in range(0, FLAGS.num_epochs):
        # Train
        print("Starting epoch {}".format(epoch))
        session.run(train_iterator_init)
        while True:
          try:
            session.run(dist_train)
          except tf.errors.OutOfRangeError:
            break
        print("Training loss: {:0.4f}, accuracy: {:0.2f}%".format(
            session.run(training_loss_result),
            session.run(training_accuracy_result) * 100))
        training_loss.reset_states()
        training_accuracy.reset_states()
        # Test
        session.run(test_iterator_init)
        while True:
          try:
            session.run(dist_test)
          except tf.errors.OutOfRangeError:
            break
        print("Test loss: {:0.4f}, accuracy: {:0.2f}%".format(
            session.run(test_loss_result),
            session.run(test_accuracy_result) * 100))
        test_loss.reset_states()
        test_accuracy.reset_states()