# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for TPUStrategy."""

from absl import logging
from absl.testing import parameterized

from tensorflow.core.protobuf import config_pb2
from tensorflow.python import tf2
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.distribute import distribute_lib
from tensorflow.python.distribute import reduce_util
from tensorflow.python.distribute import strategy_test_lib
from tensorflow.python.distribute import tpu_strategy as tpu_lib
from tensorflow.python.distribute import tpu_values
from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.eager import remote
from tensorflow.python.eager import test
from tensorflow.python.framework import composite_tensor
from tensorflow.python.framework import config
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import device as tf_device
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.framework import tensor_spec
from tensorflow.python.framework import test_util
from tensorflow.python.framework import type_spec
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_switch_case
from tensorflow.python.ops import embedding_ops
from tensorflow.python.ops import gen_dataset_ops
from tensorflow.python.ops import logging_ops
from tensorflow.python.ops import lookup_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import string_ops
from tensorflow.python.ops import variables
from tensorflow.python.ops.ragged import ragged_tensor
from tensorflow.python.platform import flags
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.tpu import device_assignment as device_assignment_lib
from tensorflow.python.tpu import tpu_hardware_feature
from tensorflow.python.tpu import tpu_replication
from tensorflow.python.tpu import tpu_strategy_util
from tensorflow.python.training import server_lib
from tensorflow.python.util import nest

FLAGS = flags.FLAGS
flags.DEFINE_string("tpu", "", "Name of TPU to connect to.")
flags.DEFINE_string("project", None, "Name of GCP project with TPU.")
flags.DEFINE_string("zone", None, "Name of GCP zone with TPU.")


def get_tpu_cluster_resolver():
  resolver = tpu_cluster_resolver.TPUClusterResolver(
      tpu=FLAGS.tpu,
      zone=FLAGS.zone,
      project=FLAGS.project,
  )
  return resolver


def get_tpu_strategy(enable_packed_var=False):
  resolver = get_tpu_cluster_resolver()
  remote.connect_to_cluster(resolver)
  tpu_cluster_resolver.initialize_tpu_system(resolver)
  strategy = tpu_lib.TPUStrategyV2(resolver)
  strategy._enable_packed_variable_in_eager_mode = enable_packed_var
  return strategy


# TPU tests which don't use TPUStrategy.
@test_util.with_eager_op_as_function
class TPUTest(test.TestCase):

  # In this case, the entire computation in foo is compiled using JIT
  # compilation.
  def test_single_tpu_jit_compile(self):
    with ops.device("/device:TPU:0"):
      a = variables.Variable(1)

    def get_a_plus_one():
      return a + 1

    @def_function.function(
        input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
    def foo(x):
      b = x + get_a_plus_one()
      b = b + get_a_plus_one()
      return b + 1

    with ops.device("/device:TPU:0"):
      result = foo(a)
    self.assertAllEqual(6, result)

  # In this case, the entire computation in foo is compiled using JIT
  # compilation and contains unsupported ops that should be outside compiled.
  def test_single_tpu_jit_compile_with_outside_compilation(self):
    context.enable_jit_compile_rewrite()
    get_tpu_strategy(True)
    config.set_soft_device_placement(True)
    with ops.device("/device:TPU:1"):
      a = variables.Variable(1)

    def get_a_plus_one():
      return a + 1

    @def_function.function(
        input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
    def foo(x):
      b = x + get_a_plus_one()
      my_str = string_ops.as_string(b)
      new_str = my_str + "0"
      c = string_ops.string_to_number(new_str, out_type=dtypes.int32)
      logging_ops.print_v2(c)
      b = c + get_a_plus_one()
      return b + 1

    with ops.device("/device:TPU:1"):
      result = foo(a)
    self.assertAllEqual(33, result)

  # In this case, each of the ops in the TPU device scope are compiled and run
  # individually.
  def test_single_tpu_on_demand(self):
    with ops.device("/device:TPU:0"):
      a = variables.Variable(1)

    def get_a_plus_one():
      return a + 1

    x = 1
    with ops.device("/device:TPU:0"):
      b = x + get_a_plus_one()
      b = b + get_a_plus_one()
    result = b + 1

    self.assertAllEqual(6, result)

  # In this case, each of the ops in the tf.function and TPU device scope are
  # compiled and run individually.
  def test_single_tpu_on_demand_tf_function(self):
    with ops.device("/device:TPU:0"):
      a = variables.Variable(1)

    def get_a_plus_one():
      return a + 1

    @def_function.function(
        input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
    def foo(x):
      with ops.device("/device:TPU:0"):
        b = x + get_a_plus_one()
        b = b + get_a_plus_one()
      return b + 1

    result = foo(a)
    self.assertAllEqual(6, result)

  def test_multiple_initialize_system(self):
    resolver = get_tpu_cluster_resolver()
    remote.connect_to_cluster(resolver)
    tpu_cluster_resolver.initialize_tpu_system(resolver)

    with test.mock.patch.object(logging, "warning") as mock_log:
      tpu_cluster_resolver.initialize_tpu_system(resolver)
      self.assertRegex(str(mock_log.call_args), "already been initialized")

  def test_initialize_tpu_system_impl_input(self):
    resolver = get_tpu_cluster_resolver()
    with self.assertRaisesRegex(
        TypeError,
        r"tpu_cluster_resolver_cls is not"
        r" tf.distribute.cluster_resolver.TPUClusterResolver."):
      tpu_strategy_util.initialize_tpu_system_impl(
          resolver, tpu_cluster_resolver_cls=None)

  def test_shutdown_tpu_system_impl_input(self):
    resolver = get_tpu_cluster_resolver()
    with self.assertRaisesRegex(
        TypeError,
        r"tpu_cluster_resolver_cls is not"
        r" tf.distribute.cluster_resolver.TPUClusterResolver."):
      tpu_strategy_util.shutdown_tpu_system_impl(
          resolver, tpu_cluster_resolver_cls=None)

  def test_tpu_tf_function_same_device(self):
    with ops.device("/device:TPU:0"):
      a = variables.Variable(1)

    @def_function.function(experimental_attributes={"_noinline": True})
    def get_a_plus_one():
      return a + 1

    @def_function.function(
        input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
    def foo(x):
      with ops.device("/device:TPU:0"):
        b = x + get_a_plus_one()
      return b + 1

    result = foo(a)
    self.assertAllEqual(4, result)

  def test_tpu_return_int32(self):
    with ops.device("/device:TPU:0"):
      a = variables.Variable(0)

    @def_function.function
    def foo():
      return a + 1

    @def_function.function
    def bar():
      with ops.device("/device:TPU:1"):
        return foo()

    with ops.device("/device:CPU:0"):
      result = bar() + 1
      self.assertAllEqual(result, 2)

  def test_tpu_output_device(self):

    def foo():
      return 1 + 1

    func1 = def_function.function(foo, jit_compile=False)
    func2 = def_function.function(
        foo,
        jit_compile=False,
        experimental_attributes={
            "_OutputsOnOpDevice": True,
        },
    )

    with ops.device("/device:TPU:0"):
      ret1 = func1()
      ret2 = func2()

    self.assertAllEqual(ret1.backing_device,
                        "/job:localhost/replica:0/task:0/device:CPU:0")
    self.assertAllEqual(ret2.backing_device,
                        "/job:localhost/replica:0/task:0/device:TPU:0")

  def test_on_demand_op_with_dynamic_output(self):
    with ops.device("/device:TPU:0"):
      where_output = array_ops.where([True, False, True])
    self.assertAllEqual(where_output, [[0], [2]])

    with ops.device("/device:TPU:0"):
      repeat_output = array_ops.repeat(math_ops.range(2), [1, 4])
    self.assertAllEqual(repeat_output, [0, 1, 1, 1, 1])


@parameterized.named_parameters([("PackedVar", True), ("", False)])
@test_util.with_eager_op_as_function
class TPUStrategyTest(test.TestCase, parameterized.TestCase):

  def test_handle_in_cross_replica_context(self, enable_packed_var):
    strategy = get_tpu_strategy(enable_packed_var)
    with strategy.scope():
      v = variables.Variable(1.0)

    @def_function.function
    def func():
      self.assertEndsWith(v.handle.device, "device:TPU:0")
      return v + 1.0

    ret = func()
    self.assertAllEqual(ret, 2.0)

  def testStaticHashTableDatasetFnHostTrainingLoop(self, enable_packed_var):
    self._dataset_fn_tracing_count = 0
    strategy = get_tpu_strategy(enable_packed_var)

    with strategy.scope():
      vals = [0, 1, 2]
      keys_tensor = constant_op.constant(
          list(range(len(vals))), dtype=dtypes.int64)
      vals_tensor = constant_op.constant(vals)
      initializer = lookup_ops.KeyValueTensorInitializer(
          keys_tensor, vals_tensor)
      per_worker_table = lookup_ops.StaticHashTable(
          initializer, default_value=-1)

    @def_function.function
    def dataset_fn(input_context):
      tensor = constant_op.constant([0, 1, 3], dtype=dtypes.int64)
      global_batch_size = 2
      batch_size = input_context.get_per_replica_batch_size(global_batch_size)
      dataset = dataset_ops.Dataset.from_tensors(tensor).repeat().batch(
          batch_size, drop_remainder=True)
      dataset = dataset.shard(input_context.num_input_pipelines,
                              input_context.input_pipeline_id)
      dataset = dataset.prefetch(2)  # This prefetches 2 batches per device.
      dataset = dataset.map(per_worker_table.lookup)
      self._dataset_fn_tracing_count += 1
      return dataset

    dist_iterator = iter(
        strategy.experimental_distribute_datasets_from_function(dataset_fn))

    @def_function.function
    def step_fn(inputs):
      # inputs should be [0, 1, -1]
      return math_ops.reduce_sum(inputs)

    def train_steps(iterator, steps):

      for _ in math_ops.range(steps):
        strategy.run(step_fn, args=(next(iterator),))

    train_steps(dist_iterator, steps=5)
    self.assertEqual(self._dataset_fn_tracing_count, 1)

  def test_function_compile_with_xla(self, enable_packed_var):
    if FLAGS.tpu_use_tfrt:
      self.skipTest(
          "This test triggers _XlaCompile and XlaLaunch which are not "
          "supported in tfrt yet. We should avoid using these kernels on TPU. "
          "However, it is a workaround to support b/129842431. We need more "
          "discussion about how to support it in the long term.")
    strategy = get_tpu_strategy(enable_packed_var)
    with strategy.scope():
      v = variables.Variable(1.0)

    @def_function.function
    def func():
      return v.read_value() + 1.0

    with ops.device("/device:TPU:0"):
      self.assertAllEqual(func(), 2.0)

  def test_sequential_runs(self, enable_packed_var):
    resolver = get_tpu_cluster_resolver()
    remote.connect_to_cluster(resolver)
    topology = tpu_cluster_resolver.initialize_tpu_system(resolver)
    # Computation replicated to all cores.
    device_assignment = device_assignment_lib.DeviceAssignment.build(
        topology, num_replicas=2)
    strategy = tpu_lib.TPUStrategyV2(
        resolver, experimental_device_assignment=device_assignment)
    strategy._enable_packed_variable_in_eager_mode = enable_packed_var

    # Computation on the 1st core.
    device_assignment2 = device_assignment_lib.DeviceAssignment.build(
        topology, num_replicas=1)
    strategy2 = tpu_lib.TPUStrategyV2(
        resolver, experimental_device_assignment=device_assignment2)

    def computation(x):
      return math_ops.square(x)

    @def_function.function
    def train_step():
      outputs = strategy.experimental_local_results(
          strategy.run(computation, args=([2., 2.],)))
      outputs2 = strategy2.run(
          computation, args=([outputs[0]],))
      return outputs2

    self.assertAllEqual([[16., 16.]], train_step())

  def test_device_switch_case(self, enable_packed_var):
    strategy = get_tpu_strategy(enable_packed_var)
    with strategy.scope():
      a = variables.Variable(1)

    inference_iteration = variables.Variable(-1)

    def inference_fn(x, i):
      return a + x + i

    @def_function.function
    def run_inference(x):

      def do_inference(device, inference_fn, i):
        with ops.device(device):
          return inference_fn(x, i)

      branch_fns = {
          0: (lambda: do_inference("/device:TPU:0", inference_fn, 0)),
          1: (lambda: do_inference("/device:TPU:1", inference_fn, 1)),
      }
      branch_index = inference_iteration.assign_add(1, use_locking=True) % 2
      return control_flow_switch_case.switch_case(branch_index, branch_fns)

    self.assertAllEqual(2., run_inference(1))  # Use TPU core 0.
    self.assertAllEqual(3., run_inference(1))  # Use TPU core 1.

  def test_recover_from_compilation_failures(self, enable_packed_var):
    # TODO(b/148150981): Stop skipping this test once recovery works
    # for non-local TPU.
    if FLAGS.tpu:
      self.skipTest("Recovery fails for non-local TPU, see b/148150981")

    # Disable automatic outside compilation.
    config.set_soft_device_placement(False)
    strategy = get_tpu_strategy(enable_packed_var)

    @def_function.function
    def compilation_failure_run():

      def computation():
        return random_ops.random_gamma([10], [0.5, 1.5])

      return strategy.run(computation)

    with self.assertRaises(errors.OpError):
      compilation_failure_run()

    @def_function.function
    def good_run():

      def computation():
        return random_ops.random_normal([10])

      return strategy.run(computation)

    good_run()

  def test_dynamic_shape_with_outside_compilation_failure(
      self, enable_packed_var):
    # Enable automatic outside compilation.
    config.set_soft_device_placement(True)
    strategy = get_tpu_strategy(enable_packed_var)
    dataset = dataset_ops.Dataset.from_tensors(("string", 1.0)).repeat().batch(
        2, drop_remainder=False)
    dataset = strategy.experimental_distribute_dataset(dataset)
    iterator = iter(dataset)

    @def_function.function
    def train_fn(iterator):

      def step_fn(inputs):
        input0, input1 = inputs
        return array_ops.size(input0), math_ops.reduce_sum(input1)

      return strategy.experimental_local_results(
          strategy.run(step_fn, args=(next(iterator),)))

    with self.assertRaises(errors.InvalidArgumentError):
      logging.info(train_fn(iterator))

  def test_computation_on_subset_cores(self, enable_packed_var):
    resolver = get_tpu_cluster_resolver()
    remote.connect_to_cluster(resolver)
    topology = tpu_cluster_resolver.initialize_tpu_system(resolver)
    all_core_strategy = tpu_lib.TPUStrategyV2(resolver)
    all_core_strategy._enable_packed_variable_in_eager_mode = enable_packed_var

    with all_core_strategy.scope():
      v = variables.Variable(0.0,
                             aggregation=variables.VariableAggregation.MEAN)

    # Computation on the 1st core.
    device_assignment = device_assignment_lib.DeviceAssignment.build(
        topology, num_replicas=1)
    first_core_strategy = tpu_lib.TPUStrategyV2(
        resolver, experimental_device_assignment=device_assignment)
    first_core_strategy._enable_packed_variable_in_eager_mode = (
        enable_packed_var)

    # Computation on the 2nd core.
    device_assignment2 = device_assignment_lib.DeviceAssignment(
        topology, [[[0, 0, 0, 1]]])
    second_core_strategy = tpu_lib.TPUStrategyV2(
        resolver, experimental_device_assignment=device_assignment2)
    second_core_strategy._enable_packed_variable_in_eager_mode = (
        enable_packed_var)

    @def_function.function
    def train_step():

      def step_fn():
        return v + 1.0

      all_core_strategy.run(step_fn)
      r1 = first_core_strategy.run(step_fn)
      r2 = second_core_strategy.run(step_fn)
      return r1 + r2

    train_step()
    self.assertAllEqual(2., train_step())

  def test_worker_devices_on_subset_cores(self, enable_packed_var):
    resolver = get_tpu_cluster_resolver()
    remote.connect_to_cluster(resolver)
    topology = tpu_cluster_resolver.initialize_tpu_system(resolver)

    # Strategy for the 1st core.
    device_assignment = device_assignment_lib.DeviceAssignment.build(
        topology, num_replicas=1)
    first_core_strategy = tpu_lib.TPUStrategyV2(
        resolver, experimental_device_assignment=device_assignment)
    first_core_strategy._enable_packed_variable_in_eager_mode = (
        enable_packed_var)

    # Strategy for the 2nd core.
    device_assignment2 = device_assignment_lib.DeviceAssignment(
        topology, [[[0, 0, 0, 1]]])
    second_core_strategy = tpu_lib.TPUStrategyV2(
        resolver, experimental_device_assignment=device_assignment2)
    second_core_strategy._enable_packed_variable_in_eager_mode = (
        enable_packed_var)

    self.assertLen(first_core_strategy.extended.worker_devices, 1)
    self.assertEndsWith(first_core_strategy.extended.worker_devices[0],
                        "device:TPU:0")

    self.assertLen(second_core_strategy.extended.worker_devices, 1)
    self.assertEndsWith(second_core_strategy.extended.worker_devices[0],
                        "device:TPU:1")

  def test_control_output_in_while_body_fn(self, enable_packed_var):
    strategy = get_tpu_strategy(enable_packed_var)

    with strategy.scope():
      v = variables.Variable(
          0.0, aggregation=variables.VariableAggregation.MEAN)

    @def_function.function
    def train_step():

      def step_fn():
        v.assign_add(1)

      for _ in math_ops.range(2):
        strategy.run(step_fn)

    train_step()
    self.assertEqual(2.0, v.numpy())

  def test_cluster_conditional_with_dynamic_shape(self, enable_packed_var):
    strategy = get_tpu_strategy(enable_packed_var)

    @def_function.function
    def train_step():

      def shape_list(tensor):
        shape = tensor.shape.as_list()

        non_static_indexes = []
        for (index, dim) in enumerate(shape):
          if dim is None:
            non_static_indexes.append(index)

        if not non_static_indexes:
          return shape

        dynamic_shape = array_ops.shape(input=tensor)
        for index in non_static_indexes:
          shape[index] = dynamic_shape[index]

        return shape

      def step_fn(condition):
        where = array_ops.where(condition)
        if array_ops.shape(where)[0] > 0:
          tensor_shape = shape_list(where)
          d1 = tensor_shape[0]
          d2 = tensor_shape[1]
          where = array_ops.reshape(where, [d1, d2])
        return where

      return strategy.run(step_fn, args=([True, False, True],))

    outputs = strategy.experimental_local_results(train_step())
    self.assertAllEqual(outputs[0].numpy(), [[0], [2]])

  def test_cluster_in_graph_and_while_body_fn(self, enable_packed_var):
    strategy = get_tpu_strategy(enable_packed_var)

    @def_function.function
    def train_step():

      def step_fn(prev):
        s = prev + 1
        return s

      def init_fn():
        return array_ops.zeros(shape=())

      prev = strategy.run(init_fn)
      for _ in math_ops.range(10):
        prev = strategy.run(step_fn, args=(prev,))
      return strategy.reduce(reduce_util.ReduceOp.SUM, prev, axis=None)

    sum_val = train_step().numpy().astype(float)
    self.assertEqual(sum_val, strategy.num_replicas_in_sync * 10)

  def test_two_clusters_with_same_fn(self, enable_packed_var):
    strategy = get_tpu_strategy(enable_packed_var)

    @def_function.function
    def foo(x):
      return strategy.run(lambda x: x + 1, (x,))

    @def_function.function
    def bar(x):
      foo(x)
      return foo(x)

    bar(1)

  def test_tpu_variable_run_argument(self, enable_packed_var):
    # TPUStrategy.run() casts inputs to Tensor, but has logic to preserve
    # variables to avoid unintuitive errors.
    # Here we test that a TPUDistributedVariable passed to TPUStrategy.run()
    # remains a variable.

    strategy = get_tpu_strategy(enable_packed_var)

    with strategy.scope():
      tpu_variable = variables.Variable(1)

    def replica_step(first_arg, variable):
      del first_arg  # Just here to make sure we're not relying on arg position.

      if variable is not None:
        self.assertIsInstance(variable, tpu_values.TPUDistributedVariable)

    @def_function.function
    def step():
      strategy.run(
          replica_step, args=(
              2,
              tpu_variable,
          ))

    step()

  def test_tpu_run_arg_parsing(self, enable_packed_var):
    strategy = get_tpu_strategy(enable_packed_var)

    with strategy.scope():
      tpu_vars = [variables.Variable(1)]

    def only_star_args(*args):
      del args

    def pos_and_star_args(first_arg, *args):
      del first_arg
      del args

    def named_args(first_arg, second_arg):
      del first_arg
      del second_arg

    def star_args_and_kw_only(*args, kw):
      del args
      del kw

    # pylint:disable=function-redefined
    @def_function.function
    def step():
      strategy.run(only_star_args, args=(2,))

    step()

    @def_function.function
    def step():
      strategy.run(named_args, kwargs={"first_arg": 2, "second_arg": 3})

    step()

    with self.assertRaisesRegex(TypeError, r"got multiple values for argument"):

      @def_function.function
      def step():
        strategy.run(
            named_args, args=(1,), kwargs={
                "first_arg": 2,
                "second_arg": 3
            })

      step()

    with self.assertRaisesRegex(ValueError,
                                r"cannot handle Variables passed to \*args"):

      @def_function.function
      def step():
        strategy.run(
            only_star_args, args=(
                2,
                tpu_vars,
            ))

      step()

    @def_function.function
    def step():
      strategy.run(pos_and_star_args, args=(2, 3, 4))

    step()

    @def_function.function
    def step():
      strategy.run(star_args_and_kw_only, args=(2, 3), kwargs={"kw": tpu_vars})

    step()

    with self.assertRaisesRegex(ValueError,
                                r"mix of positional args and \*args"):

      @def_function.function
      def step():
        strategy.run(pos_and_star_args, args=(tpu_vars, 3, 4))

      step()

    with self.assertRaisesRegex(ValueError, r"Too many positional arguments"):

      @def_function.function
      def step():
        strategy.run(named_args, args=(2, 3, 4))

      step()

    class DummyClass:

      @def_function.function
      def method(self, arg_1):
        del arg_1

      def step(self):
        strategy.run(self.method, args=(tpu_vars,))

    DummyClass().step()
    # pylint:enable=function-redefined

  def test_using_external_variable_inside_tf_function(self, enable_packed_var):
    strategy = get_tpu_strategy(enable_packed_var)
    dataset = dataset_ops.Dataset.range(
        strategy.num_replicas_in_sync * 2,
        output_type=dtypes.float32).batch(strategy.num_replicas_in_sync)
    input_iterator = iter(strategy.experimental_distribute_dataset(dataset))

    v = variables.Variable(2.0)

    @def_function.function
    def train_step(data):
      def computation(inputs):
        return inputs + v
      return strategy.run(computation, args=(data,))

    expected_result = [[x + 2.] for x in range(0, strategy.num_replicas_in_sync)
                      ]
    self.assertAllEqual(
        expected_result,
        strategy.experimental_local_results(train_step(next(input_iterator))))

  # TODO(b/145574622): Remove this test once it is re-enabled in values_test.py.
  def test_all_reduce_on_sync_on_read_variable(self, enable_packed_var):
    strategy = get_tpu_strategy(enable_packed_var)
    dataset = dataset_ops.Dataset.range(
        strategy.num_replicas_in_sync, output_type=dtypes.float32).batch(
            strategy.num_replicas_in_sync, drop_remainder=True)
    input_iterator = iter(strategy.experimental_distribute_dataset(dataset))

    with strategy.scope():
      w = variables.Variable(
          (0.,),
          shape=(1,),
          trainable=False,
          synchronization=variables.VariableSynchronization.ON_READ,
          aggregation=variables.VariableAggregation.ONLY_FIRST_REPLICA)

    self.assertFalse(w._is_mirrored())

    @def_function.function
    def run(iterator):

      def computation(x):
        w.assign(x + w)
        return w

      def all_reduce(x):
        ctx = distribute_lib.get_replica_context()
        return ctx.all_reduce("SUM", w) + x

      outputs = strategy.run(computation, args=(next(iterator),))
      outputs2 = strategy.experimental_local_results(
          strategy.run(all_reduce, args=(outputs,)))
      return outputs2

    data = range(0, strategy.num_replicas_in_sync)
    data_sum = sum(data)
    expected_result = [
        [x + data_sum] for x in range(0, strategy.num_replicas_in_sync)
    ]
    self.assertAllEqual(expected_result, run(input_iterator))
    self.assertAllEqual((0.,), w.read_value())

  def test_run_output_on_device(self, enable_packed_var):
    strategy = get_tpu_strategy(enable_packed_var)

    def computation(x):
      return math_ops.square(x)

    @def_function.function
    def train_step():
      outputs = strategy.experimental_local_results(
          strategy.run(computation, args=(2,)))
      return outputs

    results = train_step()
    self.assertAllEqual([4., 4.], results)
    self.assertAllEqual("/job:localhost/replica:0/task:0/device:TPU:0",
                        results[0].backing_device)
    self.assertAllEqual("/job:localhost/replica:0/task:0/device:TPU:1",
                        results[1].backing_device)

  def test_run_passing_and_returning_nones(self, enable_packed_var):
    strategy = get_tpu_strategy(enable_packed_var)

    @def_function.function
    def train_step():

      def computation(x):
        return x

      # Note that this input None is nested.
      outputs = strategy.experimental_local_results(
          strategy.run(computation, args=([1, [2, None]],)))
      return outputs

    results = train_step()

    self.assertAllEqual(1, results[0][0])
    self.assertAllEqual(2, results[0][1][0])
    self.assertIsNone(results[0][1][1])

  def test_run_passing_and_returning_empty_list(self, enable_packed_var):
    strategy = get_tpu_strategy(enable_packed_var)

    @def_function.function
    def train_step():

      def computation(x):
        return x

      outputs = strategy.experimental_local_results(
          strategy.run(computation, args=([],)))
      return outputs

    self.assertEqual([], train_step()[0])

  def test_run_passing_and_returning_empty_dict(self, enable_packed_var):
    strategy = get_tpu_strategy(enable_packed_var)

    @def_function.function
    def train_step():

      def computation(x):
        return x

      outputs = strategy.experimental_local_results(
          strategy.run(computation, args=({},)))
      return outputs

    self.assertEqual({}, train_step()[0])

  def test_composite_input_output(self, enable_packed_var):
    strategy = get_tpu_strategy(enable_packed_var)
    if strategy.num_replicas_in_sync != 2:
      self.skipTest("Test assumes two replicas.")

    with strategy.scope():
      table = variables.Variable(
          initial_value=[[0.0, 1.0], [3.0, 7.0]], dtype=dtypes.float32)

    @def_function.function
    def sparse_lookup(iterator):

      def tpu_function(sparse):
        # Assumes dense_shape is (2, *)
        looked_up = array_ops.gather(table, sparse.values)
        segment_sum = math_ops.unsorted_segment_sum(
            looked_up, sparse.indices[:, 0], 2)
        return sparse, segment_sum

      return nest.map_structure(
          strategy.experimental_local_results,
          strategy.run(tpu_function, args=(next(iterator),)))

    def dataset_fn(_):
      dataset = dataset_ops.Dataset.range(2)

      def make_sparse(_):
        return sparse_tensor.SparseTensor(
            indices=array_ops.constant([[0, 0], [1, 0], [1, 1]],
                                       dtype=dtypes.int64),
            values=array_ops.constant([0, 0, 1], dtype=dtypes.int32),
            dense_shape=array_ops.constant([2, 2], dtype=dtypes.int64))

      return dataset.map(make_sparse)

    dataset = iter(
        strategy.distribute_datasets_from_function(
            dataset_fn,
            distribute_lib.InputOptions(experimental_fetch_to_device=False)))

    sparse, result = sparse_lookup(dataset)

    # All replicas return identical reults.
    for replica in range(strategy.num_replicas_in_sync):
      self.assertIsInstance(sparse[replica], sparse_tensor.SparseTensor)
      self.assertAllEqual(sparse[replica].indices, [[0, 0], [1, 0], [1, 1]])
      self.assertAllEqual(sparse[replica].values, [0, 0, 1])
      self.assertAllEqual(sparse[replica].dense_shape, [2, 2])
      self.assertAllEqual(result[replica], [[0.0, 1.0], [3.0, 8.0]])

  def test_composite_input_non_flat_output(self, enable_packed_var):
    strategy = get_tpu_strategy(enable_packed_var)
    if strategy.num_replicas_in_sync != 2:
      self.skipTest("Test assumes two replicas.")

    with strategy.scope():
      table = variables.Variable(
          initial_value=[[0.0, 1.0], [3.0, 7.0]], dtype=dtypes.float32)

    @def_function.function
    def sparse_lookup(iterator):

      def tpu_function(sparse):
        # Assumes dense_shape is (2, *)
        looked_up = array_ops.gather(table, sparse.values)
        segment_sum = math_ops.unsorted_segment_sum(
            looked_up, sparse.indices[:, 0], 2)
        return {"sparse": sparse, "segment_sum": segment_sum}

      return nest.map_structure(
          strategy.experimental_local_results,
          strategy.run(tpu_function, args=(next(iterator),)))

    def dataset_fn(_):
      dataset = dataset_ops.Dataset.range(2)

      def make_sparse(_):
        return sparse_tensor.SparseTensor(
            indices=array_ops.constant([[0, 0], [1, 0], [1, 1]],
                                       dtype=dtypes.int64),
            values=array_ops.constant([0, 0, 1], dtype=dtypes.int32),
            dense_shape=array_ops.constant([2, 2], dtype=dtypes.int64))

      return dataset.map(make_sparse)

    dataset = iter(
        strategy.distribute_datasets_from_function(
            dataset_fn,
            distribute_lib.InputOptions(experimental_fetch_to_device=False)))

    output = sparse_lookup(dataset)

    # All replicas return identical reults.
    for replica in range(strategy.num_replicas_in_sync):
      self.assertIsInstance(output["sparse"][replica],
                            sparse_tensor.SparseTensor)
      self.assertAllEqual(output["sparse"][replica].indices,
                          [[0, 0], [1, 0], [1, 1]])
      self.assertAllEqual(output["sparse"][replica].values, [0, 0, 1])
      self.assertAllEqual(output["sparse"][replica].dense_shape, [2, 2])
      self.assertAllEqual(output["segment_sum"][replica],
                          [[0.0, 1.0], [3.0, 8.0]])

  def test_composite_input_dynamic_shapes_outside_compilation(
      self, enable_packed_var):
    strategy = get_tpu_strategy(enable_packed_var)
    if strategy.num_replicas_in_sync != 2:
      self.skipTest("Test assumes two replicas.")

    table = variables.Variable(
        initial_value=[[0.0, 1.0], [3.0, 7.0]], dtype=dtypes.float32)

    @def_function.function
    def sparse_lookup(iterator):

      def tpu_function(sparse):
        lookup = tpu_replication.outside_compilation(
            embedding_ops.safe_embedding_lookup_sparse, table, sparse)
        return math_ops.reduce_sum(lookup, axis=0)

      return strategy.experimental_local_results(
          strategy.run(tpu_function, args=(next(iterator),)))

    def dataset_fn(_):
      dataset = dataset_ops.Dataset.range(2)

      def make_sparse(i):
        indices = array_ops.constant([[0, 0], [1, 0], [1, 1]],
                                     dtype=dtypes.int64)[0:2 + i]
        values = array_ops.constant([0, 0, 1], dtype=dtypes.int32)[0:2 + i]
        shape = [
            array_ops.constant([2], dtype=dtypes.int64),
            array_ops.expand_dims(1 + i, axis=0)
        ]
        dense_shape = array_ops.concat(shape, axis=0)
        return sparse_tensor.SparseTensor(
            indices=indices, values=values, dense_shape=dense_shape)

      return dataset.map(make_sparse)

    dataset = iter(
        strategy.distribute_datasets_from_function(
            dataset_fn,
            options=distribute_lib.InputOptions(
                experimental_fetch_to_device=False)))

    result = sparse_lookup(dataset)
    self.assertAllEqual(result, [[0.0, 2.0], [1.5, 5.0]])

  def test_composite_input_with_non_flat_components(self, enable_packed_var):
    strategy = get_tpu_strategy(enable_packed_var)

    class TestCompositeTypeSpec(type_spec.TypeSpec):

      def __init__(self, component_type_spec):
        self._component_type_spec = component_type_spec

      @property
      def value_type(self):
        return TestComposite

      def _to_components(self, value):
        return value.values

      def _from_components(self, components):
        return TestComposite(components[0], components[1][0], components[1][1])

      @property
      def _component_specs(self):
        return [self._component_type_spec,
                [self._component_type_spec, self._component_type_spec]]

      def _serialize(self):
        return (self._component_type_spec,)

    class TestComposite(composite_tensor.CompositeTensor):

      def __init__(self, value1, value2, value3):
        self.values = [value1, [value2, value3]]

      @property
      def _type_spec(self):
        return TestCompositeTypeSpec(
            tensor_spec.TensorSpec.from_tensor(self.values[0]))

      def _shape_invariant_to_type_spec(self, shape):
        return [shape, [shape, shape]]

    @def_function.function
    def test_fn(test_composite):

      def tpu_function(composite):
        return (composite,
                composite.values[0] + (
                    composite.values[1][0] + composite.values[1][1])/2)

      return nest.map_structure(
          strategy.experimental_local_results,
          strategy.run(tpu_function, args=(test_composite,)))

    a = array_ops.constant([0.1])
    b = array_ops.constant([1.2])
    c = array_ops.constant([-0.4])
    test_composite = TestComposite(a, b, c)

    composite, result = test_fn(test_composite)

    # All replicas return identical reults.
    for replica in range(strategy.num_replicas_in_sync):
      self.assertIsInstance(composite[replica], TestComposite)
      self.assertAllEqual(composite[replica].values[0], a)
      self.assertAllEqual(composite[replica].values[1][0], b)
      self.assertAllEqual(composite[replica].values[1][1], c)
      self.assertAllEqual(result[replica], array_ops.constant([0.50000006]))

  def test_per_device_tracing_of_mirrored_variables(self, enable_packed_var):
    # Define trace_count as a list to avoid python scoping error
    trace_count = [0]

    strategy = get_tpu_strategy(enable_packed_var)
    with strategy.scope():
      variable = variables.Variable(0.0)

    @def_function.function
    def add_one():
      trace_count[0] = trace_count[0] + 1
      return math_ops.add(variable, constant_op.constant(1.0))

    @def_function.function
    def update_variable():
      for device in set(strategy.extended.worker_devices):
        with ops.device(device):
          add_one()

    with strategy.scope():
      update_variable.get_concrete_function()
      self.assertLen(strategy.extended.worker_devices, trace_count[0])

  def test_tpu_cancellation_does_not_close_chips(self, enable_packed_var):
    if not FLAGS.tpu_use_tfrt:
      self.skipTest(
          "`tpu_cancellation_closes_chip only applies to TFRT TPU Runtime.")
    strategy = get_tpu_strategy(enable_packed_var)
    num_replicas = strategy.num_replicas_in_sync
    with strategy.scope():
      x = random_ops.random_normal((10240, 10240))
      y = random_ops.random_normal((10240, 10240))

      v = variables.Variable(array_ops.identity(x))
      dist_dataset = strategy.experimental_distribute_dataset(
          dataset_ops.Dataset.from_tensors(y).repeat(num_replicas).batch(
              num_replicas))
      dist_iterator = iter(dist_dataset)

      @def_function.function
      def train_steps(v, iterator, steps):

        def step_fn(inputs):
          for val in inputs:
            v.assign(math_ops.matmul(v, val))

        for _ in math_ops.range(steps):
          strategy.run(step_fn, args=(next(iterator),))

      with self.assertRaises(errors.OutOfRangeError):
        # The iterator has num_replicas/num_replicas = 1 step only.
        train_steps(v, dist_iterator, 2)

      # If TPU chips are not closed we can run the function on TPU again.
      w = variables.Variable(array_ops.identity(x))
      dist_dataset = strategy.experimental_distribute_dataset(
          dataset_ops.Dataset.from_tensors(y).repeat(num_replicas).batch(
              num_replicas))
      dist_iterator = iter(dist_dataset)
      train_steps(w, dist_iterator, 1)

  def test_tpu_hardware_feature(self, enable_packed_var):
    strategy = get_tpu_strategy(enable_packed_var)
    self.assertIsInstance(
        strategy.extended.tpu_hardware_feature.embedding_feature,
        tpu_hardware_feature.HardwareFeature.EmbeddingFeature)

  def test_get_tpu_cluster_resolver(self, enable_packed_var):
    strategy = get_tpu_strategy(enable_packed_var)
    self.assertIsNotNone(strategy.cluster_resolver)

  def test_replica_order_for_distribute_datasets_from_function(
      self, enable_packed_var
  ):
    def _create_dataset(strategy):
      def dataset_fn(ctx):
        del ctx
        return dataset_ops.Dataset.range(2)

      return strategy.distribute_datasets_from_function(dataset_fn)

    values = self._test_replica_order(_create_dataset).values

    self.assertLen(values, 2)
    self.assertEqual(1, values[0].numpy())
    self.assertEqual(0, values[1].numpy())

  def test_replica_order_for_experimental_distribute_dataset(
      self, enable_packed_var
  ):
    def _create_dataset(strategy):
      dataset = dataset_ops.Dataset.range(2).batch(2)
      return strategy.experimental_distribute_dataset(dataset)

    values = self._test_replica_order(_create_dataset).values

    self.assertLen(values, 2)
    self.assertEqual(1, values[0].numpy())
    self.assertEqual(0, values[1].numpy())

  def _test_replica_order(self, create_dist_dataset_fn):
    tf2.enable()

    resolver = get_tpu_cluster_resolver()
    remote.connect_to_cluster(resolver)
    topology = tpu_cluster_resolver.initialize_tpu_system(resolver)
    device_assignment = device_assignment_lib.DeviceAssignment(
        topology, core_assignment=[[[0, 0, 0, 1]], [[0, 0, 0, 0]]]
    )
    strategy = tpu_lib.TPUStrategyV2(
        resolver, experimental_device_assignment=device_assignment
    )

    dist_dataset = create_dist_dataset_fn(strategy)
    iterator = iter(dist_dataset)

    @def_function.function
    def test_iterators_order(iterator):
      return next(iterator)

    return test_iterators_order(iterator)


@test_util.with_eager_op_as_function
class TPUStrategyDataPrefetchTest(test.TestCase):

  def test_prefetch_to_device_default(self):
    strategy = get_tpu_strategy()
    dataset = dataset_ops.Dataset.range(
        strategy.num_replicas_in_sync * 2,
        output_type=dtypes.float32).batch(strategy.num_replicas_in_sync)

    # Check default, should prefetch to TPU.
    dataset_item = next(iter(strategy.experimental_distribute_dataset(dataset)))
    dataset_location = tf_device.DeviceSpec.from_string(
        dataset_item.values[0].device)
    self.assertEqual(dataset_location.device_type, "TPU")

  def test_prefetch_to_device_tpu(self):
    strategy = get_tpu_strategy()
    dataset = dataset_ops.Dataset.range(
        strategy.num_replicas_in_sync * 2,
        output_type=dtypes.float32).batch(strategy.num_replicas_in_sync)

    input_options = distribute_lib.InputOptions(
        experimental_fetch_to_device=True)
    dataset_item = next(iter(strategy.experimental_distribute_dataset(
        dataset, options=input_options)))
    dataset_location = tf_device.DeviceSpec.from_string(
        dataset_item.values[0].device)
    self.assertEqual(dataset_location.device_type, "TPU")

  def test_prefetch_to_device_cpu(self):
    strategy = get_tpu_strategy()
    dataset = dataset_ops.Dataset.range(
        strategy.num_replicas_in_sync * 2,
        output_type=dtypes.float32).batch(strategy.num_replicas_in_sync)

    # Should be CPU when prefetch_to_device is False.
    input_options = distribute_lib.InputOptions(
        experimental_fetch_to_device=False)
    dataset_item = next(iter(strategy.experimental_distribute_dataset(
        dataset, options=input_options)))
    dataset_location = tf_device.DeviceSpec.from_string(
        dataset_item.values[0].device)
    self.assertEqual(dataset_location.device_type, "CPU")

  def test_prefetch_to_device_sparse_dataset(self):
    strategy = get_tpu_strategy()
    # Values here aren't important.
    dataset = dataset_ops.Dataset.from_tensors(
        sparse_tensor.SparseTensor(indices=[[0, 0], [0, 1], [1, 0]],
                                   values=[1, 2, 3],
                                   dense_shape=[2, 2]))
    dataset = dataset.repeat()
    dataset = dataset.batch(strategy.num_replicas_in_sync)

    with self.assertRaisesRegex(ValueError, "TPUStrategy does not support"):
      iter(strategy.experimental_distribute_dataset(dataset))

  def test_prefetch_to_device_ragged_dataset(self):
    strategy = get_tpu_strategy()
    # Values here aren't important.
    dataset = dataset_ops.Dataset.from_tensors(
        ragged_tensor.RaggedTensor.from_row_splits(
            values=[1, 2, 3],
            row_splits=[0, 2, 3]))
    dataset = dataset.repeat()
    dataset = dataset.batch(strategy.num_replicas_in_sync)

    with self.assertRaisesRegex(ValueError, "TPUStrategy does not support"):
      iter(strategy.experimental_distribute_dataset(dataset))

  def test_prefetch_to_device_sparse_dataset_fn(self):
    strategy = get_tpu_strategy()
    def dataset_fn(ctx):
      del ctx
      # Values here aren't important.
      dataset = dataset_ops.Dataset.from_tensors(
          sparse_tensor.SparseTensor(indices=[[0, 0], [0, 1], [1, 0]],
                                     values=[1, 2, 3],
                                     dense_shape=[2, 2]))
      dataset = dataset.repeat()
      return dataset.batch(strategy.num_replicas_in_sync)

    with self.assertRaisesRegex(ValueError, "TPUStrategy does not support"):
      iter(strategy.distribute_datasets_from_function(dataset_fn))

  def test_prefetch_to_device_ragged_dataset_fn(self):
    strategy = get_tpu_strategy()
    def dataset_fn(ctx):
      del ctx
      # Values here aren't important.
      dataset = dataset_ops.Dataset.from_tensors(
          ragged_tensor.RaggedTensor.from_row_splits(
              values=[1, 2, 3],
              row_splits=[0, 2, 3]))
      dataset = dataset.repeat()
      return dataset.batch(strategy.num_replicas_in_sync)

    with self.assertRaisesRegex(ValueError, "TPUStrategy does not support"):
      iter(strategy.distribute_datasets_from_function(dataset_fn))

  def test_create_iterator_on_device(self):

    @def_function.function
    def create_iter():
      with ops.device("/device:TPU:0"):
        return gen_dataset_ops.anonymous_iterator_v3(
            output_types=[dtypes.float32], output_shapes=[[]])

    create_iter()


@test_util.with_eager_op_as_function
class TPUStrategyDistributionTest(
    strategy_test_lib.DistributionTestBase,
    strategy_test_lib.TwoDeviceDistributionTestBase):

  def test_update_config_proto(self):
    resolver = get_tpu_cluster_resolver()
    remote.connect_to_cluster(resolver)
    tpu_cluster_resolver.initialize_tpu_system(resolver)
    strategy = tpu_lib.TPUStrategyV2(resolver)

    config_proto = config_pb2.ConfigProto()
    cluster_spec = server_lib.ClusterSpec({"worker": ["fake1", "fake2"]})
    with test.mock.patch.object(
        resolver, "cluster_spec", return_value=cluster_spec):
      new_config = strategy.update_config_proto(config_proto)

    # Verify cluster_def.
    self.assertProtoEquals(cluster_spec.as_cluster_def(),
                           new_config.cluster_def)

    # Verify isolate_session_state
    self.assertTrue(new_config.isolate_session_state)

  def test_make_input_fn_iterable(self):
    dataset_fn = lambda: dataset_ops.Dataset.range(10)
    expected_values = [[i, i+1] for i in range(0, 10, 2)]
    distribution = get_tpu_strategy()
    input_fn = self._input_fn_to_test_input_context(
        dataset_fn,
        expected_num_replicas_in_sync=2,
        expected_num_input_pipelines=1,
        expected_input_pipeline_id=0)
    self._test_input_fn_iterable(distribution, input_fn, expected_values)

  def test_make_input_fn_iterator(self):
    dataset_fn = lambda: dataset_ops.Dataset.range(10)
    expected_values = [[i, i+1] for i in range(0, 10, 2)]
    distribution = get_tpu_strategy()
    input_fn = self._input_fn_to_test_input_context(
        dataset_fn,
        expected_num_replicas_in_sync=2,
        expected_num_input_pipelines=1,
        expected_input_pipeline_id=0)
    iterator = distribution.make_input_fn_iterator(input_fn)
    self._test_input_fn_iterator(
        iterator,
        distribution.extended.worker_devices,
        expected_values)

  def test_num_replicas_in_sync(self):
    strategy = get_tpu_strategy()
    self.assertEqual(2, strategy.num_replicas_in_sync)

  def test_call_and_merge_exceptions(self):
    strategy = get_tpu_strategy()
    self._test_call_and_merge_exceptions(strategy)

  def test_numpy_dataset(self):
    strategy = get_tpu_strategy()
    self._test_numpy_dataset(strategy, run_in_function=True)

  def test_global_step_update(self):
    strategy = get_tpu_strategy()
    self._test_global_step_update(strategy)

  def test_run(self):
    strategy = get_tpu_strategy()
    self._test_run(strategy, run_in_function=True)

  def test_summary_for_replica_zero_only(self):
    strategy = get_tpu_strategy()
    self._test_summary_for_replica_zero_only(strategy)

  def test_all_reduce_sum(self):
    strategy = get_tpu_strategy()
    self._test_all_reduce_sum(strategy, run_in_function=True)

  def test_all_reduce_sum_gradients(self):
    strategy = get_tpu_strategy()
    self._test_all_reduce_sum_gradients(strategy, run_in_function=True)

  def test_all_reduce_sum_gradient_tape(self):
    strategy = get_tpu_strategy()
    self._test_all_reduce_sum_gradient_tape(strategy, run_in_function=True)

  def test_all_reduce_mean(self):
    strategy = get_tpu_strategy()
    self._test_all_reduce_mean(strategy, run_in_function=True)

  def test_all_reduce_mean_gradients(self):
    strategy = get_tpu_strategy()
    self._test_all_reduce_mean_gradients(strategy, run_in_function=True)

  def test_all_reduce_mean_gradient_tape(self):
    strategy = get_tpu_strategy()
    self._test_all_reduce_mean_gradient_tape(strategy, run_in_function=True)

  def test_reduce(self):
    strategy = get_tpu_strategy()

    inputs = strategy.make_input_fn_iterator(
        lambda _: dataset_ops.Dataset.from_tensor_slices([2., 3.]))

    self.evaluate(inputs.initialize())
    per_replica_outputs = strategy.run(
        def_function.function(math_ops.square), args=(next(inputs),))

    with strategy.scope():
      mean = strategy.reduce(reduce_util.ReduceOp.MEAN, per_replica_outputs,
                             axis=None)
      self.assertEqual(6.5, self.evaluate(mean))

  def test_constraint(self):
    strategy = get_tpu_strategy()

    with strategy.scope():
      variable = variables.Variable(initial_value=2.,
                                    constraint=lambda x: 0. * x + 1.)
    self.assertEqual(variable.value().numpy(), 2)

    @def_function.function
    def update_variable():
      variable.assign_add(1)
      variable.assign(variable.constraint(variable))

    update_variable()
    self.assertEqual(variable.value().numpy(), 1)

  def test_trainable_variables(self):
    strategy = get_tpu_strategy()
    self._test_trainable_variable(strategy)


@test_util.with_eager_op_as_function
class DeviceAssignmentTest(test.TestCase):

  def test_core_assignment(self):
    resolver = get_tpu_cluster_resolver()
    remote.connect_to_cluster(resolver)
    topology = tpu_cluster_resolver.initialize_tpu_system(resolver)
    device_assignment = device_assignment_lib.DeviceAssignment(
        topology, core_assignment=[[[0, 0, 0, 0]]])
    self.assertAllEqual([[[0, 0, 0, 0]]], device_assignment.core_assignment)
    self.assertEqual(1, device_assignment.num_cores_per_replica)
    self.assertEqual(1, device_assignment.num_replicas)
    self.assertEqual("/task:0/device:TPU:0", device_assignment.tpu_device())
    self.assertEqual("/task:0/device:CPU:0", device_assignment.host_device())

  def test_device_assignment_strategy_properties(self):
    resolver = get_tpu_cluster_resolver()
    remote.connect_to_cluster(resolver)
    topology = tpu_cluster_resolver.initialize_tpu_system(resolver)
    device_assignment = device_assignment_lib.DeviceAssignment(
        topology, core_assignment=[[[0, 0, 0, 0]]])
    strategy = tpu_lib.TPUStrategyV2(
        resolver,
        experimental_device_assignment=device_assignment)
    self.assertEqual(strategy.extended.num_hosts, 1)
    self.assertEqual(strategy.num_replicas_in_sync, 1)
    self.assertEqual(strategy.extended.num_replicas_per_host, 1)  # pylint: disable=protected-access

  def test_device_assignment_constants(self):
    resolver = get_tpu_cluster_resolver()
    remote.connect_to_cluster(resolver)
    topology = tpu_cluster_resolver.initialize_tpu_system(resolver)
    device_assignment = device_assignment_lib.DeviceAssignment(
        topology,
        core_assignment=device_assignment_lib.SINGLE_CORE_ASSIGNMENT)
    self.assertAllEqual([[[0, 0, 0, 0]]], device_assignment.core_assignment)
    self.assertEqual(1, device_assignment.num_cores_per_replica)
    self.assertEqual(1, device_assignment.num_replicas)
    self.assertEqual("/task:0/device:TPU:0", device_assignment.tpu_device())
    self.assertEqual("/task:0/device:CPU:0", device_assignment.host_device())

  def test_variables_mismatched_device_assignment(self):
    resolver = get_tpu_cluster_resolver()
    remote.connect_to_cluster(resolver)
    topology = tpu_cluster_resolver.initialize_tpu_system(resolver)

    strategy0 = tpu_lib.TPUStrategyV2(resolver)
    self.assertEqual(
        ("/job:localhost/replica:0/task:0/device:TPU:0",
         "/job:localhost/replica:0/task:0/device:TPU:1"),
        strategy0.extended.worker_devices)

    with strategy0.scope():
      v = variables.Variable(1.)

    v1_assign_op = strategy0.experimental_local_results(v)[1].assign(42.)

    with self.cached_session():
      self.evaluate(variables.global_variables_initializer())
      self.evaluate(v1_assign_op)
      self.assertAllEqual([1., 42.],
                          self.evaluate(
                              strategy0.experimental_local_results(v)))

    # Second strategy has devices reversed relative to the first.
    device_assignment = device_assignment_lib.DeviceAssignment(
        topology, core_assignment=[[[0, 0, 0, 1]], [[0, 0, 0, 0]]])
    strategy1 = tpu_lib.TPUStrategyV2(
        resolver,
        experimental_device_assignment=device_assignment)
    self.assertEqual(
        ("/job:localhost/replica:0/task:0/device:TPU:1",
         "/job:localhost/replica:0/task:0/device:TPU:0"),
        strategy1.extended.worker_devices)

    v_read = strategy1.run(def_function.function(v.read_value))

    with self.cached_session():
      self.assertAllEqual([42., 1.],
                          self.evaluate(
                              strategy0.experimental_local_results(v_read)))


if __name__ == "__main__":
  test.main()