Suppress unrelated Cholesky decomposition error in Hypothesis test of batch shapes.

PiperOrigin-RevId: 386462611

Author: axch

tensorflow_probability/python/distributions/distribution_properties_test.py

@@ -372,7 +372,8 @@ def testCanConstructAndSampleDistribution(self, data):
   def testInferredBatchShapeMatchesTrueBatchShape(self, dist_name, data):
     dist = data.draw(
         dhps.distributions(dist_name=dist_name, validate_args=False))
-    lp = dist.log_prob(dist.sample(seed=test_util.test_seed()))
+    with tfp_hps.no_cholesky_decomposition_errors():
+      lp = dist.log_prob(dist.sample(seed=test_util.test_seed()))
 
     self.assertAllEqual(dist.batch_shape_tensor(), tf.shape(lp))
     self.assertAllEqual(dist.batch_shape, tf.shape(lp))

tensorflow_probability/python/internal/hypothesis_testlib.py

@@ -763,6 +763,24 @@ def no_tf_rank_errors():
       raise
 
 
+@contextlib.contextmanager
+def no_cholesky_decomposition_errors():
+  # Use this to suppress Cholesky errors if needed in tests where the
+  # numerics are beside the point.
+  pat = ('Cholesky decomposition was not successful. '
+         'The input might not be valid. [Op:Cholesky]')
+  try:
+    yield
+  except tf.errors.InvalidArgumentError as e:
+    msg = str(e)
+    if pat in msg:
+      # Tried an input regime where a Cholesky decomposition failed
+      # and crashed the program.
+      hp.assume(False)
+    else:
+      raise
+
+
 @contextlib.contextmanager
 def finite_ground_truth_only():
   # Recognizing the error message from python/internal/numerics_testing.py