message: Re-enable Model.predict for TPU Pods.
diff: @@ -74,12 +74,9 @@ def disable_multi_worker(method):
  """Decorator that disallows multi-worker use of `method`."""
  def _method_wrapper(self, *args, **kwargs):
  - strategy = self.distribute_strategy
  - if (self._in_multi_worker_mode() or dist_utils.is_tpu_strategy(strategy) and # pylint: disable=protected-access
  - strategy.extended.num_hosts > 1):
  + if self._in_multi_worker_mode(): # pylint: disable=protected-access
  raise ValueError('{} is not supported in multi-worker mode.'.format(
  method.__name__))
  -
  return method(self, *args, **kwargs)
  return tf_decorator.make_decorator(
  @@ -1632,9 +1629,12 @@ def reduce_per_replica(values, strategy, reduction='first'):
  if not isinstance(v, ds_values.PerReplica):
  return v
  elif reduction == 'first':
  - return strategy.unwrap(v)[0] # pylint: disable=protected-access
  + return strategy.unwrap(v)[0]
  elif reduction == 'concat':
  - return concat(strategy.unwrap(v)) # pylint: disable=protected-access
  + if _is_tpu_multi_host(strategy):
  + return _tpu_multi_host_concat(v, strategy)
  + else:
  + return concat(strategy.unwrap(v))
  else:
  raise ValueError('`reduction` must be "first" or "concat".')
  @@ -1659,3 +1659,23 @@ def to_numpy(tensors):
  return t # Don't turn ragged or sparse tensors to NumPy.
  return nest.map_structure(_to_single_numpy, tensors)
  +
  +
  +def _is_tpu_multi_host(strategy):
  + return (dist_utils.is_tpu_strategy(strategy) and
  + strategy.extended.num_hosts > 1)
  +
  +
  +def _tpu_multi_host_concat(v, strategy):
  + """Correctly order TPU PerReplica objects."""
  + replicas = strategy.unwrap(v)
  + # When distributed datasets are created from Tensors / NumPy,
  + # TPUStrategy.experimental_distribute_dataset shards data in
  + # (Replica, Host) order, and TPUStrategy.unwrap returns it in
  + # (Host, Replica) order.
  + # TODO(b/150317897): Figure out long-term plan here.
  + num_replicas_per_host = strategy.extended.num_replicas_per_host
  + ordered_replicas = []
  + for replica_id in range(num_replicas_per_host):
  + ordered_replicas += replicas[replica_id::num_replicas_per_host]
  + return concat(ordered_replicas)
  