from tensorflow.python.framework import ops

    ps = server_lib.Server(cs, job_name="ps", start=True)

    return worker, ps

  @contextlib.contextmanager
  def _maybeWithDevice(self, device):
    if device is not None:
      with ops.device(device):
        yield
    else:
      yield

  def _setupDense(self, is_distributed, dtype):
    with self._maybeWithDevice("/job:ps" if is_distributed else None):
