message: Remove the workaround for collective cancellation issue in keras
  It's no longer needed
diff: @@ -2771,11 +2771,7 @@ def _collective_all_reduce_multi_worker(strategy):
  # for all strategies
  def _multi_worker_concat(v, strategy):
  """Order PerReplica objects for CollectiveAllReduceStrategy and concat."""
  - replicas = strategy.gather(v, axis=0) # pylint: disable=protected-access
  - # TODO(b/170435030): We now need to make sure these run after the iterator
  - # GetNext, so that we don't trigger aborting collective ops in the case of
  - # EOF. Remove after the issue is fixed.
  - with ops.control_dependencies([replicas]):
  + replicas = strategy.gather(v, axis=0)
  # v might not have the same shape on different replicas
  if isinstance(v, ds_values.PerReplica):
  shapes = array_ops.concat([
  @@ -2787,8 +2783,7 @@ def _multi_worker_concat(v, strategy):
  else:
  # v is a tensor. This may happen when, say, we have 2x1 multi-worker.
  all_shapes = strategy.gather(
  - array_ops.expand_dims_v2(array_ops.shape(v)[0], axis=0),
  - axis=0)
  + array_ops.expand_dims_v2(array_ops.shape(v)[0], axis=0), axis=0)
  replicas = array_ops.split(
  replicas,
  