import tensorflow.contrib.eager as tfe

  def tearDown(self):
    shutil.rmtree(self._tmp_logdir)
    super(LinearRegressionTest, self).tearDown()

  def testSyntheticDataset(self):
    true_w = tf.random_uniform([3, 1])
    true_b = [1.0]
    batch_size = 10
    num_batches = 2
    noise_level = 0.
    dataset = linear_regression.synthetic_dataset(true_w, true_b, noise_level,
                                                  batch_size, num_batches)

    it = tfe.Iterator(dataset)
    for _ in range(2):
      (xs, ys) = it.next()
      self.assertEqual((batch_size, 3), xs.shape)
      self.assertEqual((batch_size, 1), ys.shape)
      self.assertEqual(tf.float32, xs.dtype)
      self.assertEqual(tf.float32, ys.dtype)
    with self.assertRaises(StopIteration):
      it.next()

  def testLinearRegression(self):
    true_w = [[1.0], [-0.5], [2.0]]
    true_b = [1.0]
