from tensorflow.python.ops import variables

      yield

  def _setupDense(self, is_distributed, dtype):
    with self._maybeWithDevice("/job:ps" if is_distributed else None):
      var0 = variables.Variable([[0.0, 1.0], [2.0, 3.0]], dtype=dtype)
      var1 = variables.Variable([4.0, 5.0], dtype=dtype)
    with self._maybeWithDevice("/job:worker" if is_distributed else None):
      grads0 = constant_op.constant([[0.1, 0.1], [0.1, 0.1]], dtype=dtype)
