from tensorflow.python.training import gradient_descent

        yield
    else:
      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)
      grads1 = constant_op.constant([0.01, 0.01], dtype=dtype)
      sgd = gradient_descent.GradientDescentOptimizer(3.0)
      clip_opt = variable_clipping_optimizer.VariableClippingOptimizer(
          sgd, {var0: [1]}, 2.0)

      update_op = clip_opt.apply_gradients(
          list(zip([grads0, grads1], [var0, var1])))
      variables.global_variables_initializer().run()
    return var0, var1, update_op

  def _assertDenseCorrect(self, var0, var1, update_op):
    # Fetch params to validate initial values
