Enable numpy testing for distributions / bijectors.

srvasude · tensorflower-gardener · commit a7b35afc857b · 2022-06-16T19:00:59.000-07:00
PiperOrigin-RevId: 455513834
diff --git a/tensorflow_probability/python/bijectors/BUILD b/tensorflow_probability/python/bijectors/BUILD
@@ -1254,7 +1254,6 @@ multi_substrate_py_test(
     size = "medium",
     srcs = ["correlation_cholesky_test.py"],
     jax_size = "large",
-    numpy_tags = ["notap"],
     shard_count = 2,
     deps = [
         ":bijector_test_util",
@@ -1636,7 +1635,6 @@ multi_substrate_py_test(
     name = "permute_test",
     size = "small",
     srcs = ["permute_test.py"],
-    numpy_tags = ["notap"],
     deps = [
         ":bijector_test_util",
         ":bijectors",
@@ -1705,7 +1703,6 @@ multi_substrate_py_test(
     name = "rational_quadratic_spline_test",
     size = "medium",
     srcs = ["rational_quadratic_spline_test.py"],
-    numpy_tags = ["notap"],
     tags = ["hypothesis"],
     deps = [
         ":bijector_test_util",
diff --git a/tensorflow_probability/python/bijectors/bijector_test.py b/tensorflow_probability/python/bijectors/bijector_test.py
@@ -904,7 +904,8 @@ def testInstanceCache(self):
     z = instance_cache_bijector.forward(x)
     self.assertIsNot(y, z)
 
-  @test_util.jax_disable_test_missing_functionality('keras')
+  @test_util.disable_test_for_backend(
+      disable_numpy=True, disable_jax=True, reason='keras')
   @parameterized.named_parameters(
       ('Keras', True),
       ('NoKeras', False))
@@ -1093,13 +1094,15 @@ def test_caches(self):
 @test_util.test_all_tf_execution_regimes
 class TfModuleTest(test_util.TestCase):
 
+  @test_util.numpy_disable_variable_test
   @test_util.jax_disable_variable_test
   def test_variable_tracking(self):
     x = tf.Variable(1.)
     b = ForwardOnlyBijector(scale=x, validate_args=True)
     self.assertIsInstance(b, tf.Module)
     self.assertEqual((x,), b.trainable_variables)
 
+  @test_util.numpy_disable_variable_test
   @test_util.jax_disable_variable_test
   def test_gradient(self):
     x = tf.Variable(1.)
diff --git a/tensorflow_probability/python/bijectors/correlation_cholesky_test.py b/tensorflow_probability/python/bijectors/correlation_cholesky_test.py
@@ -227,6 +227,7 @@ def sample_mcmc_chain():
               cdf=beta_dist.cdf,
               false_fail_rate=1e-9))
 
+  @test_util.numpy_disable_gradient_test
   def testTheoreticalFldj(self):
     bijector = tfb.CorrelationCholesky()
     x = np.linspace(-50, 50, num=30).reshape(5, 6).astype(np.float64)
diff --git a/tensorflow_probability/python/bijectors/permute_test.py b/tensorflow_probability/python/bijectors/permute_test.py
@@ -83,8 +83,9 @@ def testBijectiveAndFiniteAxis(self):
         bijector, x, y, eval_func=self.evaluate, event_ndims=2, rtol=1e-6,
         atol=0)
 
-  @test_util.jax_disable_test_missing_functionality(
-      'Test specific to Keras with losing shape information.')
+  @test_util.disable_test_for_backend(
+      disable_numpy=True, disable_jax=True,
+      reason='Test specific to Keras with losing shape information.')
   def testPreservesShape(self):
     # TODO(b/131157549, b/131124359): Test should not be needed. Consider
     # deleting when underlying issue with constant eager tensors is fixed.
diff --git a/tensorflow_probability/python/bijectors/rational_quadratic_spline_test.py b/tensorflow_probability/python/bijectors/rational_quadratic_spline_test.py
@@ -144,6 +144,7 @@ def testDegenerateSplines(self):
         np.zeros_like(xs),
         self.evaluate(bijector.forward_log_det_jacobian(xs, event_ndims=0)))
 
+  @test_util.numpy_disable_gradient_test
   def testTheoreticalFldjSimple(self):
     bijector = tfb.RationalQuadraticSpline(
         bin_widths=[1., 1],
diff --git a/tensorflow_probability/python/distributions/BUILD b/tensorflow_probability/python/distributions/BUILD
@@ -2668,7 +2668,6 @@ multi_substrate_py_test(
     size = "medium",
     srcs = ["distribution_test.py"],
     jax_size = "medium",
-    numpy_tags = ["notap"],
     deps = [
         # numpy dep,
         # tensorflow dep,
@@ -2918,7 +2917,6 @@ multi_substrate_py_test(
 multi_substrate_py_test(
     name = "half_cauchy_test",
     srcs = ["half_cauchy_test.py"],
-    numpy_tags = ["notap"],
     shard_count = 5,
     deps = [
         # numpy dep,
@@ -2932,7 +2930,6 @@ multi_substrate_py_test(
 multi_substrate_py_test(
     name = "half_normal_test",
     srcs = ["half_normal_test.py"],
-    numpy_tags = ["notap"],
     deps = [
         # numpy dep,
         # scipy dep,
@@ -3157,7 +3154,6 @@ multi_substrate_py_test(
 multi_substrate_py_test(
     name = "lambertw_f_test",
     srcs = ["lambertw_f_test.py"],
-    numpy_tags = ["notap"],
     deps = [
         ":lambertw_f",
         # absl/testing:parameterized dep,
@@ -3228,7 +3224,6 @@ multi_substrate_py_test(
 multi_substrate_py_test(
     name = "loglogistic_test",
     srcs = ["loglogistic_test.py"],
-    numpy_tags = ["notap"],
     deps = [
         # numpy dep,
         # scipy dep,
@@ -3339,7 +3334,6 @@ multi_substrate_py_test(
     size = "medium",
     srcs = ["mixture_same_family_test.py"],
     jax_size = "large",
-    numpy_tags = ["notap"],
     shard_count = 3,
     deps = [
         # hypothesis dep,
@@ -3370,7 +3364,6 @@ multi_substrate_py_test(
     size = "medium",
     srcs = ["multinomial_test.py"],
     jax_size = "medium",
-    numpy_tags = ["notap"],
     shard_count = 5,
     deps = [
         # hypothesis dep,
@@ -3426,7 +3419,6 @@ multi_substrate_py_test(
 multi_substrate_py_test(
     name = "mvn_diag_test",
     srcs = ["mvn_diag_test.py"],
-    numpy_tags = ["notap"],
     deps = [
         # numpy dep,
         # scipy dep,
@@ -3454,7 +3446,6 @@ multi_substrate_py_test(
 multi_substrate_py_test(
     name = "mvn_linear_operator_test",
     srcs = ["mvn_linear_operator_test.py"],
-    numpy_tags = ["notap"],
     deps = [
         # numpy dep,
         # scipy dep,
@@ -3481,7 +3472,6 @@ multi_substrate_py_test(
     name = "mvn_tril_test",
     size = "medium",
     srcs = ["mvn_tril_test.py"],
-    numpy_tags = ["notap"],
     shard_count = 7,
     tags = ["nomsan"],
     deps = [
diff --git a/tensorflow_probability/python/distributions/distribution_test.py b/tensorflow_probability/python/distributions/distribution_test.py
@@ -787,6 +787,7 @@ def __init__(self, extra_param, *args, **kwargs):
 @test_util.test_all_tf_execution_regimes
 class TfModuleTest(test_util.TestCase):
 
+  @test_util.numpy_disable_variable_test
   @test_util.jax_disable_variable_test
   def test_variable_tracking_works(self):
     scale = tf.Variable(1.)
diff --git a/tensorflow_probability/python/distributions/half_cauchy_test.py b/tensorflow_probability/python/distributions/half_cauchy_test.py
@@ -411,6 +411,7 @@ def testHalfCauchySampleMultidimensionalMedian(self):
     expected_shape = (tf.TensorShape(n_).concatenate(half_cauchy.batch_shape))
     self.assertAllEqual(expected_shape, samples.shape)
 
+  @test_util.numpy_disable_gradient_test
   def testHalfCauchyPdfGradientZeroOutsideSupport(self):
     loc_ = [-3.1, -2., 0., 1.1]
     loc = self._create_placeholder_with_default(loc_, name='loc')
@@ -430,6 +431,7 @@ def testHalfCauchyPdfGradientZeroOutsideSupport(self):
         grads,
         [np.zeros_like(loc_), 0., np.zeros_like(loc_)])
 
+  @test_util.numpy_disable_gradient_test
   def testHalfCauchyCdfGradientZeroOutsideSupport(self):
     loc_ = [-3.1, -2., 0., 1.1]
     loc = self._create_placeholder_with_default(loc_, name='loc')
@@ -449,6 +451,7 @@ def testHalfCauchyCdfGradientZeroOutsideSupport(self):
         grads,
         [np.zeros_like(loc_), 0., np.zeros_like(loc_)])
 
+  @test_util.numpy_disable_gradient_test
   def testHalfCauchyGradientsAndValueFiniteAtLoc(self):
     batch_size = 1000
     loc_ = np.linspace(0., 100., batch_size)
@@ -470,6 +473,7 @@ def testHalfCauchyGradientsAndValueFiniteAtLoc(self):
       for grad in grads:
         self.assertAllFinite(grad)
 
+  @test_util.numpy_disable_gradient_test
   def testHalfCauchyGradientsAndValueFiniteAtGreaterThanLoc(self):
     batch_size = 1000
     loc = self._create_placeholder_with_default([0.] * batch_size, name='loc')
diff --git a/tensorflow_probability/python/distributions/half_normal_test.py b/tensorflow_probability/python/distributions/half_normal_test.py
@@ -154,6 +154,7 @@ def testHalfNormalQuantile(self):
     expected_x = sp_stats.halfnorm(scale=scale).ppf(p)
     self.assertAllClose(expected_x, self.evaluate(x), atol=0)
 
+  @test_util.numpy_disable_gradient_test
   @parameterized.parameters(np.float32, np.float64)
   def testFiniteGradients(self, dtype):
     scale = tf.constant(dtype(3.0))
@@ -318,12 +319,12 @@ def testHalfNormalHalfNormalKL(self):
 
     kl = tfd.kl_divergence(a, b)
 
-    x = a.sample(int(4e5), seed=test_util.test_seed(hardcoded_seed=0))
+    x = a.sample(int(6e5), seed=test_util.test_seed(hardcoded_seed=0))
     kl_samples = a.log_prob(x) - b.log_prob(x)
 
     kl_, kl_samples_ = self.evaluate([kl, kl_samples])
     self.assertAllClose(kl_, true_kl, atol=2e-15)
-    self.assertAllMeansClose(kl_samples_, true_kl, axis=0, atol=0., rtol=5e-2)
+    self.assertAllMeansClose(kl_samples_, true_kl, axis=0, atol=0., rtol=8e-2)
 
     zero_kl = tfd.kl_divergence(a, a)
     true_zero_kl_, zero_kl_ = self.evaluate([tf.zeros_like(zero_kl), zero_kl])
diff --git a/tensorflow_probability/python/distributions/johnson_su_test.py b/tensorflow_probability/python/distributions/johnson_su_test.py
@@ -219,6 +219,7 @@ def testJohnsonSULogCDF(self):
     expected_cdf = sp_stats.johnsonsu.logcdf(x, skewness, tailweight, mu, sigma)
     self.assertAllClose(expected_cdf, self.evaluate(cdf), atol=0, rtol=1e-3)
 
+  @test_util.numpy_disable_gradient_test
   def testFiniteGradientAtDifficultPoints(self):
     def make_fn(dtype, attr):
       x = np.array([-100., -20., -5., 0., 5., 20., 100.]).astype(dtype)
@@ -318,9 +319,11 @@ def _testQuantileFiniteGradientAtDifficultPoints(self, dtype):
     self.assertAllFinite(grads[0])
     self.assertAllFinite(grads[1])
 
+  @test_util.numpy_disable_gradient_test
   def testQuantileFiniteGradientAtDifficultPointsFloat32(self):
     self._testQuantileFiniteGradientAtDifficultPoints(np.float32)
 
+  @test_util.numpy_disable_gradient_test
   def testQuantileFiniteGradientAtDifficultPointsFloat64(self):
     self._testQuantileFiniteGradientAtDifficultPoints(np.float64)
 
@@ -389,6 +392,7 @@ def testJohnsonSUSample(self):
     self.assertAllEqual(expected_samples_shape, samples.shape)
     self.assertAllEqual(expected_samples_shape, sample_values.shape)
 
+  @test_util.numpy_disable_gradient_test
   def testJohnsonSUFullyReparameterized(self):
     skewness = tf.constant(1.0)
     tailweight = tf.constant(2.0)
diff --git a/tensorflow_probability/python/distributions/lambertw_f_test.py b/tensorflow_probability/python/distributions/lambertw_f_test.py
@@ -169,8 +169,9 @@ def testVariance(self, delta, variance_multiplier):
     else:
       self.assertAllClose(lwn.variance(), scale**2 * variance_multiplier)
 
-  @test_util.jax_disable_test_missing_functionality(
-      "Keras/Layers are not supported in JAX.")
+  @test_util.disable_test_for_backend(
+      disable_numpy=True, disable_jax=True,
+      reason="Keras/Layers are not supported in JAX or Numpy.")
   def testWorksInDistributionLayerAndNegloglik(self):
     """Test that distribution works as layer and in gradient optimization."""
     x = np.random.uniform(size=(100, 1))
diff --git a/tensorflow_probability/python/distributions/loglogistic.py b/tensorflow_probability/python/distributions/loglogistic.py
@@ -167,5 +167,5 @@ def sinc(x, name=None):
   """Calculate the (normalized) sinus cardinalis of x."""
   name = name or 'sinc'
   with tf.name_scope(name):
-    x *= np.pi
+    x = np.pi * x
     return tf.where(x != 0., tf.math.sin(x) / x, 1.)
diff --git a/tensorflow_probability/python/distributions/mixture_same_family_test.py b/tensorflow_probability/python/distributions/mixture_same_family_test.py
@@ -119,7 +119,7 @@ def testLogCdf(self):
         axis=1)
     actual_log_cdf_, expected_log_cdf_ = self.evaluate(
         [actual_log_cdf, expected_log_cdf])
-    self.assertAllClose(actual_log_cdf_, expected_log_cdf_, rtol=1e-6, atol=0.0)
+    self.assertAllClose(actual_log_cdf_, expected_log_cdf_, rtol=2e-5, atol=0.0)
 
   def testCovarianceWithBatch(self):
     d = self._build_mvndiag_mixture(
@@ -247,6 +247,7 @@ def testReparameterizationOfNonReparameterizedComponents(self):
           reparameterize=True,
           validate_args=True)
 
+  @test_util.numpy_disable_gradient_test
   def testSecondGradientIsDisabled(self):
     if not self.use_static_shape:
       return
@@ -296,6 +297,7 @@ def exact(*parameters):
     _, expected = tfp.math.value_and_gradient(exact, parameters)
     self.assertAllClose(actual, expected, atol=0.1, rtol=0.2)
 
+  @test_util.numpy_disable_gradient_test
   def testReparameterizationGradientsNormalScalarComponents(self):
     def mixture_func(logits, loc, scale):
       return tfd.MixtureSameFamily(
@@ -313,6 +315,7 @@ def mixture_func(logits, loc, scale):
           function,
           num_samples=10000)
 
+  @test_util.numpy_disable_gradient_test
   def testReparameterizationGradientsNormalVectorComponents(self):
     def mixture_func(logits, loc, scale):
       return tfd.MixtureSameFamily(
@@ -331,6 +334,7 @@ def mixture_func(logits, loc, scale):
           function,
           num_samples=20000)
 
+  @test_util.numpy_disable_gradient_test
   def testReparameterizationGradientsNormalMatrixComponents(self):
     def mixture_func(logits, loc, scale):
       return tfd.MixtureSameFamily(
@@ -350,6 +354,7 @@ def mixture_func(logits, loc, scale):
           function,
           num_samples=50000)
 
+  @test_util.numpy_disable_gradient_test
   def testReparameterizationGradientsExponentialScalarComponents(self):
     def mixture_func(logits, rate):
       return tfd.MixtureSameFamily(
@@ -462,6 +467,7 @@ def testExcessiveConcretizationOfParams(self):
     # additional call to `_parameter_control_dependencies`, and thus an
     # additional concretizations of the underlying distribution parameters.
 
+  @test_util.numpy_disable_gradient_test
   def testExcessiveConcretizationOfParamsWithReparameterization(self):
     logits = tfp_hps.defer_and_count_usage(self._build_variable(
         np.zeros(5), name='logits', static_rank=True))
diff --git a/tensorflow_probability/python/distributions/multinomial_test.py b/tensorflow_probability/python/distributions/multinomial_test.py
@@ -444,6 +444,7 @@ def testParamTensorFromProbs(self):
 @test_util.test_all_tf_execution_regimes
 class MultinomialFromVariableTest(test_util.TestCase):
 
+  @test_util.numpy_disable_gradient_test
   @test_util.jax_disable_variable_test
   def testGradientLogits(self):
     x = tf.Variable([-1., 0., 1])
@@ -454,6 +455,7 @@ def testGradientLogits(self):
     self.assertLen(g, 1)
     self.assertAllNotNone(g)
 
+  @test_util.numpy_disable_gradient_test
   @test_util.jax_disable_variable_test
   def testGradientProbs(self):
     x = tf.Variable([0.1, 0.7, 0.2])
diff --git a/tensorflow_probability/python/distributions/multivariate_student_t_test.py b/tensorflow_probability/python/distributions/multivariate_student_t_test.py
@@ -340,6 +340,7 @@ def testSamplingDeterministic(self):
     samples2 = self.evaluate(dist2.sample(100, seed=seed))
     self.assertAllClose(samples1, samples2)
 
+  @test_util.numpy_disable_gradient_test
   def testSamplingFullyReparameterized(self):
     df = self._input(2.)
     loc = self._input([1., 2.])
diff --git a/tensorflow_probability/python/distributions/mvn_diag_test.py b/tensorflow_probability/python/distributions/mvn_diag_test.py
@@ -230,6 +230,7 @@ def testStddev(self, is_static):
     self.assertAllClose(
         np.array([[3., 2, 1], [4, 5, 6]]), self.evaluate(mvn.stddev()))
 
+  @test_util.numpy_disable_gradient_test
   def testMultivariateNormalDiagNegLogLikelihood(self):
     num_draws = 50
     dims = 3
@@ -287,6 +288,7 @@ def testDynamicEventShape(self):
     self.assertListEqual(tensorshape_util.as_list(mvn.batch_shape), [2, 3])
     self.assertListEqual(tensorshape_util.as_list(mvn.event_shape), [None])
 
+  @test_util.numpy_disable_gradient_test
   def testKLDivIdenticalGradientDefined(self):
     dims = 3
     loc = tf.zeros([dims], dtype=tf.float32)
@@ -348,6 +350,7 @@ def testVariableScaleIdentityMultiplier(self):
       lp = d.log_prob([0., 0.])
     self.assertIsNotNone(tape.gradient(lp, scale_identity_multiplier))
 
+  @test_util.numpy_disable_variable_test
   @test_util.jax_disable_variable_test
   def testVariableScaleDiagAssertions(self):
     # We test that changing the scale to be non-invertible raises an exception
@@ -363,6 +366,7 @@ def testVariableScaleDiagAssertions(self):
       with tf.control_dependencies([scale_diag.assign([1., 0.])]):
         self.evaluate(d.sample(seed=test_util.test_seed()))
 
+  @test_util.numpy_disable_variable_test
   @test_util.jax_disable_variable_test
   def testVariableScaleIdentityMultiplierAssertions(self):
     # We test that changing the scale to be non-invertible raises an exception
diff --git a/tensorflow_probability/python/distributions/mvn_linear_operator_test.py b/tensorflow_probability/python/distributions/mvn_linear_operator_test.py
@@ -160,6 +160,7 @@ def testVariableLocation(self):
       lp = d.log_prob([0., 0.])
     self.assertIsNotNone(tape.gradient(lp, loc))
 
+  @test_util.numpy_disable_variable_test
   @test_util.jax_disable_variable_test
   def testVariableScaleAssertions(self):
     # We test that changing the scale to be non-invertible raises an exception
diff --git a/tensorflow_probability/python/distributions/mvn_tril_test.py b/tensorflow_probability/python/distributions/mvn_tril_test.py
diff --git a/tensorflow_probability/python/distributions/normal_test.py b/tensorflow_probability/python/distributions/normal_test.py
diff --git a/tensorflow_probability/python/distributions/probit_bernoulli_test.py b/tensorflow_probability/python/distributions/probit_bernoulli_test.py
diff --git a/tensorflow_probability/python/distributions/von_mises_test.py b/tensorflow_probability/python/distributions/von_mises_test.py
diff --git a/tensorflow_probability/python/math/BUILD b/tensorflow_probability/python/math/BUILD
diff --git a/tensorflow_probability/python/math/scan_associative_test.py b/tensorflow_probability/python/math/scan_associative_test.py
