Merge branch 'tensorflow:main' into betaincgrad

Author: leandrolcampos

spinoffs/oryx/oryx/core/trace_util.py

@@ -17,7 +17,6 @@
 import threading
 from typing import Any, Dict, Generator, List
 
-import jax
 from jax import abstract_arrays
 from jax import api_util
 from jax import core as jax_core
@@ -63,8 +62,6 @@ def wrapped(*args, **kwargs):
     flat_args, in_tree = tree_util.tree_flatten(args)
     flat_fun, out_tree = api_util.flatten_fun_nokwargs(fun, in_tree)
     flat_avals = safe_map(get_shaped_aval, flat_args)
-    if not jax.config.omnistaging_enabled:
-      raise ValueError('Oryx must be used with JAX omnistaging enabled.')
     if dynamic:
       jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(
           flat_fun,

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",

tensorflow_probability/python/bijectors/bijector_properties_test.py

@@ -32,7 +32,6 @@
 from tensorflow_probability.python.internal import tensor_util
 from tensorflow_probability.python.internal import tensorshape_util
 from tensorflow_probability.python.internal import test_util
-from tensorflow_probability.python.util.deferred_tensor import DeferredTensor
 
 from tensorflow.python.util import nest  # pylint: disable=g-direct-tensorflow-import
 
@@ -789,13 +788,6 @@ def testCompositeTensor(self, bijector_name, data):
                         'bijectors.')
       self.skipTest('`_Invert` bijectors are not `CompositeTensor`s.')
 
-    if not tf.executing_eagerly():
-      bijector = tf.nest.map_structure(
-          lambda x: (tf.convert_to_tensor(x)  # pylint: disable=g-long-lambda
-                     if isinstance(x, DeferredTensor) else x),
-          bijector,
-          expand_composites=True)
-
     self.assertIsInstance(bijector, tf.__internal__.CompositeTensor)
     flat = tf.nest.flatten(bijector, expand_composites=True)
     unflat = tf.nest.pack_sequence_as(bijector, flat, expand_composites=True)
@@ -833,6 +825,8 @@ def testCompositeTensor(self, bijector_name, data):
     grads = tape.gradient(ys, wrt_vars)
     assert_no_none_grad(bijector, 'forward', wrt_vars, grads)
 
+    self.assertConvertVariablesToTensorsWorks(bijector)
+
 
 def ensure_nonzero(x):
   return tf.where(x < 1e-6, tf.constant(1e-3, x.dtype), x)

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.)

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)

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.

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],

tensorflow_probability/python/distributions/BUILD

@@ -36,7 +36,6 @@ multi_substrate_py_library(
     srcs = ["__init__.py"],
     substrates_omit_deps = [
         ":pixel_cnn",
-        ":mixture",
     ],
     deps = [
         ":autoregressive",
@@ -2669,7 +2668,6 @@ multi_substrate_py_test(
     size = "medium",
     srcs = ["distribution_test.py"],
     jax_size = "medium",
-    numpy_tags = ["notap"],
     deps = [
         # numpy dep,
         # tensorflow dep,
@@ -2919,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,
@@ -2933,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,
@@ -3158,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,
@@ -3229,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,
@@ -3322,7 +3316,6 @@ multi_substrate_py_test(
     name = "mixture_test",
     size = "medium",
     srcs = ["mixture_test.py"],
-    jax_tags = ["notap"],
     numpy_tags = ["notap"],
     deps = [
         # hypothesis dep,
@@ -3341,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,
@@ -3372,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,
@@ -3428,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,
@@ -3456,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,
@@ -3483,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 = [

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.)

tensorflow_probability/python/distributions/gaussian_process.py

@@ -38,6 +38,8 @@
 from tensorflow_probability.python.internal import tensorshape_util
 from tensorflow_probability.python.math.psd_kernels.internal import util as psd_kernels_util
 from tensorflow.python.util import deprecation  # pylint: disable=g-direct-tensorflow-import
+from tensorflow.python.util import variable_utils  # pylint: disable=g-direct-tensorflow-import
+
 
 __all__ = [
     'GaussianProcess',
@@ -767,6 +769,14 @@ def _type_spec(self):
         omit_kwargs=('parameters', '_check_marginal_cholesky_fn'),
         non_identifying_kwargs=('name',))
 
+  def _convert_variables_to_tensors(self):
+    # pylint: disable=protected-access
+    components = self._type_spec._to_components(self)
+    tensor_components = variable_utils.convert_variables_to_tensors(
+        components)
+    return self._type_spec._from_components(tensor_components)
+    # pylint: enable=protected-access
+
 
 @auto_composite_tensor.type_spec_register(
     'tfp.distributions.GaussianProcess_ACTTypeSpec')