def _num_workers(distribution_strategy):
  """Returns number of workers in TPU pod or 1 for TPU donut."""
  # Input pipeline logit is slightly different for TPU pod and donut.
  # This function is a helper to determine number of devices to distribute
  # input pipeline on.
  # For TPU donut (2x2 topology) input pipeline is automatically distributed
  # by distribution strategy, thus thus function returns one.
  # For larger slices of TPU pod, input pipeline is manually distributed across
  # number of devices returned by this function.
  is_tpu_pod = distribution_strategy.extended._input_workers.num_workers > 1  # pylint: disable=protected-access
  if is_tpu_pod:
    return len(distribution_strategy.extended.worker_devices)
  else:
    return 1