import tensorflow as tf

      optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.1)
      linear_regression.fit(model, dataset, optimizer, logdir=self._tmp_logdir)

      self.assertAllClose(true_w, model.variables[0].numpy(), rtol=1e-2)
      self.assertAllClose(true_b, model.variables[1].numpy(), rtol=1e-2)
      self.assertTrue(glob.glob(os.path.join(self._tmp_logdir, "events.out.*")))


class EagerLinearRegressionBenchmark(tf.test.Benchmark):

  def benchmarkEagerLinearRegression(self):
    num_epochs = 10
    num_batches = 200
    batch_size = 64
    dataset = linear_regression.synthetic_dataset(
        w=tf.random_uniform([3, 1]),
        b=tf.random_uniform([1]),
        noise_level=0.01,
        batch_size=batch_size,
        num_batches=num_batches)
    burn_in_dataset = dataset.take(10)

    model = linear_regression.LinearModel()

    with tf.device(device()):
      optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.1)

      # Perform burn-in.
      linear_regression.fit(model, burn_in_dataset, optimizer)

