Fix failing test in dataset_ops_test caused by TensorFlow changing an exception type.

PiperOrigin-RevId: 215014489

research/astronet/astronet/ops/dataset_ops_test.py

@@ -48,11 +48,6 @@ class DatasetOpsTest(tf.test.TestCase):
       self.assertEqual([5], tensor_1d.shape)
       self.assertAllEqual([0, 1, 2, 3, 4], tensor_1d.eval())
 
-      # Invalid to pad Tensor with batch size 5 to batch size 3.
-      tensor_1d_pad3 = dataset_ops.pad_tensor_to_batch_size(tensor_1d, 3)
-      with self.assertRaises(tf.errors.InvalidArgumentError):
-        tensor_1d_pad3.eval()
-
       tensor_1d_pad5 = dataset_ops.pad_tensor_to_batch_size(tensor_1d, 5)
       self.assertEqual([5], tensor_1d_pad5.shape)
       self.assertAllEqual([0, 1, 2, 3, 4], tensor_1d_pad5.eval())
@@ -66,11 +61,6 @@ class DatasetOpsTest(tf.test.TestCase):
       self.assertEqual([3, 3], tensor_2d.shape)
       self.assertAllEqual([[0, 1, 2], [3, 4, 5], [6, 7, 8]], tensor_2d.eval())
 
-      tensor_2d_pad2 = dataset_ops.pad_tensor_to_batch_size(tensor_2d, 2)
-      # Invalid to pad Tensor with batch size 2 to batch size 2.
-      with self.assertRaises(tf.errors.InvalidArgumentError):
-        tensor_2d_pad2.eval()
-
       tensor_2d_pad3 = dataset_ops.pad_tensor_to_batch_size(tensor_2d, 3)
       self.assertEqual([3, 3], tensor_2d_pad3.shape)
       self.assertAllEqual([[0, 1, 2], [3, 4, 5], [6, 7, 8]],

research/astronet/astrowavenet/data/BUILD

+package(default_visibility = ["//visibility:public"])
+
+licenses(["notice"])  # Apache 2.0
+
+py_library(
+    name = "synthetic_transit_maker",
+    srcs = [
+        "synthetic_transit_maker.py",
+    ],
+)
+
+py_test(
+    name = "synthetic_transit_maker_test",
+    srcs = ["synthetic_transit_maker_test.py"],
+    srcs_version = "PY2AND3",
+    deps = [":synthetic_transit_maker"],
+)

research/astronet/astrowavenet/data/synthetic_transit_maker_test.py

@@ -92,8 +92,8 @@ class SyntheticTransitMakerTest(absltest.TestCase):
     time = np.linspace(0, 100, 100)
 
     flux, mask = transit_maker.random_light_curve(time)
-    self.assertAllClose(flux, gold_flux)
-    self.assertAllClose(mask, np.ones(100))
+    np.testing.assert_array_almost_equal(flux, gold_flux)
+    np.testing.assert_array_almost_equal(mask, np.ones(100))
 
   def testRandomLightCurveGenerator(self):
     transit_maker = synthetic_transit_maker.SyntheticTransitMaker()