from tensorflow.contrib.eager.python.examples.linear_regression import linear_regression


  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)
