Deflake linear_gaussian_ssm_test and log_normal_test

PiperOrigin-RevId: 205160259
tensorflow · Jul 18, 2018 · 4e98ee2 · 4e98ee2
1 parent c069ac5
commit 4e98ee2
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/tensorflow_probability/python/distributions/linear_gaussian_ssm_test.py b/tensorflow_probability/python/distributions/linear_gaussian_ssm_test.py
@@ -518,7 +518,8 @@ def testTooManyDimsThrowsError(self):
 
     full_batch_shape, dist = self.build_inputs([5, 4, 2, 3], [6, 5, 4, 2, 3])
 
-    with self.assertRaisesError("Cannot broadcast"):
+    with self.assertRaisesError(
+        "(Broadcasting is not supported|Cannot broadcast)"):
       self.maybe_evaluate(
           _augment_sample_shape(dist, full_batch_shape,
                                 validate_args=True))

diff --git a/tensorflow_probability/python/distributions/lognormal_test.py b/tensorflow_probability/python/distributions/lognormal_test.py
@@ -55,7 +55,7 @@ def testLogNormalStats(self):
   def testLogNormalSample(self):
     loc, scale = 1.5, 0.4
     dist = tfd.LogNormal(loc=loc, scale=scale)
-    samples = self.evaluate(dist.sample(5000))
+    samples = self.evaluate(dist.sample(6000, seed=1234))
     self.assertAllClose(np.mean(samples),
                         self.evaluate(dist.mean()),
                         atol=0.1)