message: Use device_assignment in TPUStrategy.
diff: @@ -148,7 +148,7 @@ class TPUStrategy(distribute_lib.Strategy):
  Args:
  tpu_cluster_resolver: A tf.distribute.cluster_resolver.TPUClusterResolver,
  which provides information about the TPU cluster.
  - device_assignment: Optional `tf.contrib.tpu.DeviceAssignment` to specify
  + device_assignment: Optional `tpu.experimental.DeviceAssignment` to specify
  the placement of replicas on the TPU cluster. Currently only supports
  the usecase of using a single core within a TPU cluster.
  """
  @@ -182,7 +182,7 @@ class TPUStrategyV1(distribute_lib.StrategyV1):
  metrics, summaries etc.
  This parameter is only used when Distribution Strategy is used with
  estimator or keras.
  - device_assignment: Optional `tf.contrib.tpu.DeviceAssignment` to specify
  + device_assignment: Optional `tpu.experimental.DeviceAssignment` to specify
  the placement of replicas on the TPU cluster. Currently only supports
  the usecase of using a single core within a TPU cluster.
  """
  @@ -363,7 +363,8 @@ class TPUExtended(distribute_lib.StrategyExtendedV1):
  replicate_inputs.append((nest.map_structure(
  select_replica, per_replica_inputs),))
  - replicate_outputs = tpu.replicate(run_fn, replicate_inputs)
  + replicate_outputs = tpu.replicate(
  + run_fn, replicate_inputs, device_assignment=self._device_assignment)
  # If run_fn has tensor outputs, tpu.replicate returns a list of list. We
  # will flatten it in this case. If run_fn has no tensor outputs,
  @@ -426,7 +427,7 @@ class TPUExtended(distribute_lib.StrategyExtendedV1):
  """Experimental method added to be used by Estimator.
  This is a private method only to be used by Estimator. Other frameworks
  - should directly be calling `tf.contrib.distribute.initialize_tpu_system`
  + should directly be calling `tpu.experimental.initialize_tpu_system`
  """
  tpu_strategy_util.initialize_tpu_system(self._tpu_cluster_resolver)
  @@ -720,7 +721,10 @@ class TPUExtended(distribute_lib.StrategyExtendedV1):
  maximum_shapes = None
  with strategy.scope():
  - replicate_outputs = tpu.replicate(replicated_fn, replicate_inputs,
  + replicate_outputs = tpu.replicate(
  + replicated_fn,
  + replicate_inputs,
  + device_assignment=self._device_assignment,
  maximum_shapes=maximum_shapes)
  # Remove all no ops that may have been added during 'tpu.replicate()'
  