message: TPUStrategy: Correctly set devices when device_assignment is set.
diff: @@ -38,6 +38,7 @@ from tensorflow.python.distribute.cluster_resolver import TPUClusterResolver
  from tensorflow.python.eager import context
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import constant_op
  +from tensorflow.python.framework import device_spec
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_shape
  @@ -184,15 +185,22 @@ class TPUExtended(distribute_lib.StrategyExtendedV1):
  self._tpu_metadata = get_tpu_system_metadata(self._tpu_cluster_resolver)
  self._device_assignment = device_assignment
  - # TODO(jhseu): Switch to DeviceAssignment to support pods and model
  - # parallelism.
  self._tpu_devices = [d.name for d in self._tpu_metadata.devices
  if "device:TPU:" in d.name]
  + # Only create variables for the number of replicas we're running.
  + if device_assignment is not None:
  + job_name = device_spec.DeviceSpecV2.from_string(self._tpu_devices[0]).job
  +
  + self._tpu_devices = []
  + for replica_id in range(device_assignment.num_replicas):
  + tpu_device = device_assignment.tpu_device(
  + replica=replica_id, logical_core=0, job=job_name)
  + tpu_device = device_util.canonicalize(tpu_device)
  + self._tpu_devices.append(tpu_device)
  +
  self._host_device = device_util.get_host_for_device(self._tpu_devices[0])
  - # Only create variables for the number of replicas we're running.
  - self._tpu_devices = self._tpu_devices[:self._num_replicas_in_sync]
  self._device_map = values.ReplicaDeviceMap(self._tpu_devices)
  # Preload the data onto the TPUs.
  