def ps_devices(all_workers=False):
    """List of ps devices (where to put the experts).

    Args:
      all_workers: whether the list is for all async workers or just this one.

    Returns:
      a list of device names
    """
    if ps_replicas > 0:
      if ps_gpu > 0:
        return [
            ps_job + "/task:%d/GPU:%d" % (d, gpu)
            for (d, gpu) in _ps_gpus(all_workers=all_workers)
        ]
      else:
        return [
            ps_job + "/task:%d" % d
            for d in _ps_replicas(all_workers=all_workers)
        ]
    else:
      if worker_gpu > 0:
        return ["gpu:%d" % d for d in _gpu_order(worker_gpu)]
      else:
        return [""]
List of ps devices (where to put the experts).

    Args:
      all_workers: whether the list is for all async workers or just this one.

    Returns:
      a list of device names
    
tensor2tensor/utils/devices.py