Update in find_root_secant in order to use dynamic tensor shape R0.13

tensorflow_probability/python/__init__.py

@@ -55,7 +55,7 @@ def _validate_tf_environment(package):
   #
   # Update this whenever we need to depend on a newer TensorFlow release.
   #
-  required_tensorflow_version = '2.4'
+  required_tensorflow_version = '2.5'
 #   required_tensorflow_version = '1.15'  # Needed internally -- DisableOnExport
 
   if (distutils.version.LooseVersion(tf.__version__) <

tensorflow_probability/python/bijectors/absolute_value.py

@@ -22,16 +22,14 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 
 __all__ = [
     'AbsoluteValue',
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class AbsoluteValue(bijector.AutoCompositeTensorBijector):
   """Computes `Y = g(X) = Abs(X)`, element-wise.
 

tensorflow_probability/python/bijectors/ascending.py

@@ -22,16 +22,14 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 
 
 __all__ = [
     'Ascending',
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Ascending(bijector.AutoCompositeTensorBijector):
   """Maps unconstrained R^n to R^n in ascending order.
 

tensorflow_probability/python/bijectors/bijector.py

@@ -20,6 +20,7 @@
 
 import abc
 import contextlib
+# import functools
 
 # Dependency imports
 import numpy as np
@@ -1598,26 +1599,39 @@ def _composite_tensor_shape_params(self):
       return ()
 
 
-class AutoCompositeTensorBijector(
-    Bijector, auto_composite_tensor.AutoCompositeTensor):
-  r"""Base for `CompositeTensor` bijectors with auto-generated `TypeSpec`s.
+# Temporarily disable AutoCT for TFP 0.13 release
+# class AutoCompositeTensorBijector(
+#     Bijector, auto_composite_tensor.AutoCompositeTensor):
+#   r"""Base for `CompositeTensor` bijectors with auto-generated `TypeSpec`s.
 
-  `CompositeTensor` objects are able to pass in and out of `tf.function` and
-  `tf.while_loop`, or serve as part of the signature of a TF saved model.
-  `Bijector` subclasses that follow the contract of
-  `tfp.experimental.auto_composite_tensor` may be defined as `CompositeTensor`s
-  by inheriting from `AutoCompositeTensorBijector` and applying a class
-  decorator as shown here:
+#   `CompositeTensor` objects are able to pass in and out of `tf.function` and
+#   `tf.while_loop`, or serve as part of the signature of a TF saved model.
+#   `Bijector` subclasses that follow the contract of
+#   `tfp.experimental.auto_composite_tensor` may be defined as `CompositeTensor`s  # pylint: disable=line-too-long
+#   by inheriting from `AutoCompositeTensorBijector` and applying a class
+#   decorator as shown here:
 
-  ```python
-  @tfp.experimental.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='my_module')
-  class MyBijector(tfb.AutoCompositeTensorBijector):
+#   ```python
+#   @tfp.experimental.auto_composite_tensor(
+#     omit_kwargs=('name',), module_name='my_module')
+#   class MyBijector(tfb.AutoCompositeTensorBijector):
 
-    # The remainder of the subclass implementation is unchanged.
-  ```
-  """
-  pass
+#     # The remainder of the subclass implementation is unchanged.
+#   ```
+#   """
+#   pass
+
+
+# auto_composite_tensor_bijector = functools.partial(
+#     auto_composite_tensor.auto_composite_tensor,
+#     omit_kwargs=('parameters',),
+#     non_identifying_kwargs=('name',),
+#     module_name='tfp.bijectors')
+
+AutoCompositeTensorBijector = Bijector
+
+
+auto_composite_tensor_bijector = lambda cls, **kwargs: cls
 
 
 def check_valid_ndims(ndims, validate=True):

tensorflow_probability/python/bijectors/bijector_properties_test.py

@@ -34,7 +34,7 @@
 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_probability.python.util.deferred_tensor import DeferredTensor
 
 
 TF2_FRIENDLY_BIJECTORS = (
@@ -185,6 +185,7 @@
 
 COMPOSITE_TENSOR_IS_BROKEN = [
     'BatchNormalization',  # tf.layers arg
+    'Inline',  # callable
     'RationalQuadraticSpline',  # TODO(b/185628453): Debug loss of static info.
 ]
 
@@ -197,7 +198,6 @@
 # TODO(b/182603117): Enable AutoCT for meta-bijectors and LinearOperator.
 AUTO_COMPOSITE_TENSOR_IS_BROKEN = [
     'FillScaleTriL',
-    'Invert',
     'ScaleMatvecDiag',
     'ScaleMatvecTriL',
 ]
@@ -244,8 +244,8 @@ def _constraint(param):
 
 
 # TODO(b/141098791): Eliminate this.
-@experimental.auto_composite_tensor
-class CallableModule(tf.Module, experimental.AutoCompositeTensor):
+# @experimental.auto_composite_tensor
+class CallableModule(tf.Module):  # , experimental.AutoCompositeTensor):
   """Convenience object for capturing variables closed over by Inline."""
 
   def __init__(self, fn, varobj):
@@ -630,7 +630,7 @@ def exception(bijector):
         return False
       if isinstance(bijector, tfb.Softfloor):
         return True
-      if isinstance(bijector, tfb.Invert):
+      if is_invert(bijector):
         return exception(bijector.bijector)
       return False
     if (bijector.forward_min_event_ndims == 0 and
@@ -887,32 +887,16 @@ def testEquality(self, bijector_name, data):
   @hp.given(hps.data())
   @tfp_hps.tfp_hp_settings()
   def testCompositeTensor(self, bijector_name, data):
-
+    # Test that making a composite tensor of this bijector doesn't throw any
+    # errors.
     bijector, event_dim = self._draw_bijector(
-        bijector_name, data,
-        batch_shape=[],
-        validate_args=True,
+        bijector_name, data, batch_shape=[],
         allowed_bijectors=(set(TF2_FRIENDLY_BIJECTORS) -
                            set(COMPOSITE_TENSOR_IS_BROKEN)))
-
-    # TODO(b/182603117): Remove "if" condition and s/composite_bij/bijector
-    # when AutoCT is enabled for meta-bijectors and LinearOperator.
-    if type(bijector).__name__ in AUTO_COMPOSITE_TENSOR_IS_BROKEN:
-      composite_bij = experimental.as_composite(bijector)
-    else:
-      composite_bij = bijector
-
-    if not tf.executing_eagerly():
-      composite_bij = tf.nest.map_structure(
-          lambda x: (tf.convert_to_tensor(x)  # pylint: disable=g-long-lambda
-                     if isinstance(x, DeferredTensor) else x),
-          composite_bij,
-          expand_composites=True)
-
-    self.assertIsInstance(composite_bij, tf.__internal__.CompositeTensor)
+    composite_bij = experimental.as_composite(bijector)
     flat = tf.nest.flatten(composite_bij, expand_composites=True)
-    unflat = tf.nest.pack_sequence_as(
-        composite_bij, flat, expand_composites=True)
+    unflat = tf.nest.pack_sequence_as(composite_bij, flat,
+                                      expand_composites=True)
 
     # Compare forward maps before and after compositing.
     n = 3
@@ -927,26 +911,6 @@ def testCompositeTensor(self, bijector_name, data):
     after_xs = unflat.inverse(ys)
     self.assertAllClose(*self.evaluate((before_xs, after_xs)))
 
-    # Input to tf.function
-    self.assertAllClose(
-        before_ys,
-        tf.function(lambda b: b.forward(xs))(composite_bij),
-        rtol=COMPOSITE_TENSOR_RTOL[bijector_name],
-        atol=COMPOSITE_TENSOR_ATOL[bijector_name])
-
-    # Forward mapping: Check differentiation through forward mapping with
-    # respect to the input and parameter variables.  Also check that any
-    # variables are not referenced overmuch.
-    xs = self._draw_domain_tensor(bijector, data, event_dim)
-    wrt_vars = [xs] + [v for v in composite_bij.trainable_variables
-                       if v.dtype.is_floating]
-    with tf.GradientTape() as tape:
-      tape.watch(wrt_vars)
-      # TODO(b/73073515): Fix graph mode gradients with bijector caching.
-      ys = bijector.forward(xs + 0)
-    grads = tape.gradient(ys, wrt_vars)
-    assert_no_none_grad(bijector, 'forward', wrt_vars, grads)
-
 
 def ensure_nonzero(x):
   return tf.where(x < 1e-6, tf.constant(1e-3, x.dtype), x)

tensorflow_probability/python/bijectors/bijector_test.py

@@ -25,7 +25,7 @@
 import tensorflow.compat.v1 as tf1
 import tensorflow.compat.v2 as tf
 from tensorflow_probability.python import bijectors as tfb
-from tensorflow_probability.python.internal import auto_composite_tensor
+from tensorflow_probability.python.bijectors import bijector as bijector_lib
 from tensorflow_probability.python.internal import cache_util
 from tensorflow_probability.python.internal import tensor_util
 from tensorflow_probability.python.internal import test_util
@@ -768,7 +768,7 @@ def testNestedCondition(self):
       mock_method.assert_called_once_with(mock.ANY, arg1=arg1, arg2=arg2)
 
 
-@auto_composite_tensor.auto_composite_tensor(omit_kwargs=('name',))
+@bijector_lib.auto_composite_tensor_bijector
 class CompositeForwardBijector(tfb.AutoCompositeTensorBijector):
 
   def __init__(self, scale=2., validate_args=False, name=None):
@@ -790,25 +790,26 @@ def _forward_log_det_jacobian(self, _):
     return tf.math.log(self._scale)
 
 
-@test_util.test_all_tf_execution_regimes
-class AutoCompositeTensorBijectorTest(test_util.TestCase):
+  # Test disabled temporarily for TFP 0.13 release.
+# @test_util.test_all_tf_execution_regimes
+# class AutoCompositeTensorBijectorTest(test_util.TestCase):
 
-  def test_disable_ct_bijector(self):
+#   def test_disable_ct_bijector(self):
 
-    ct_bijector = CompositeForwardBijector()
-    self.assertIsInstance(ct_bijector, tf.__internal__.CompositeTensor)
+#     ct_bijector = CompositeForwardBijector()
+#     self.assertIsInstance(ct_bijector, tf.__internal__.CompositeTensor)
 
-    non_ct_bijector = ForwardOnlyBijector()
-    self.assertNotIsInstance(non_ct_bijector, tf.__internal__.CompositeTensor)
+#     non_ct_bijector = ForwardOnlyBijector()
+#     self.assertNotIsInstance(non_ct_bijector, tf.__internal__.CompositeTensor)
 
-    flat = tf.nest.flatten(ct_bijector, expand_composites=True)
-    unflat = tf.nest.pack_sequence_as(
-        ct_bijector, flat, expand_composites=True)
+#     flat = tf.nest.flatten(ct_bijector, expand_composites=True)
+#     unflat = tf.nest.pack_sequence_as(
+#         ct_bijector, flat, expand_composites=True)
 
-    x = tf.constant([2., 3.])
-    self.assertAllClose(
-        non_ct_bijector.forward(x),
-        tf.function(lambda b: b.forward(x))(unflat))
+#     x = tf.constant([2., 3.])
+#     self.assertAllClose(
+#         non_ct_bijector.forward(x),
+#         tf.function(lambda b: b.forward(x))(unflat))
 
 
 if __name__ == '__main__':

tensorflow_probability/python/bijectors/categorical_to_discrete.py

@@ -36,7 +36,8 @@
 ]
 
 
-class CategoricalToDiscrete(bijector.Bijector):
+@bijector.auto_composite_tensor_bijector
+class CategoricalToDiscrete(bijector.AutoCompositeTensorBijector):
   """Bijector which computes `Y = g(X) = values[X]`.
 
   Example Usage:

tensorflow_probability/python/bijectors/cholesky_outer_product.py

@@ -24,7 +24,6 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import distribution_util
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import prefer_static as ps
@@ -36,8 +35,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class CholeskyOuterProduct(bijector.AutoCompositeTensorBijector):
   """Compute `g(X) = X @ X.T`; X is lower-triangular, positive-diagonal matrix.
 

tensorflow_probability/python/bijectors/cholesky_to_inv_cholesky.py

@@ -23,7 +23,6 @@
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.bijectors.cholesky_outer_product import CholeskyOuterProduct
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import prefer_static as ps
 
@@ -33,8 +32,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class CholeskyToInvCholesky(bijector.AutoCompositeTensorBijector):
   """Maps the Cholesky factor of `M` to the Cholesky factor of `M^{-1}`.
 

tensorflow_probability/python/bijectors/correlation_cholesky.py

@@ -24,7 +24,6 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.bijectors import fill_triangular
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import prefer_static as ps
 from tensorflow_probability.python.internal import tensorshape_util
 
@@ -33,8 +32,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class CorrelationCholesky(bijector.AutoCompositeTensorBijector):
   """Maps unconstrained reals to Cholesky-space correlation matrices.
 

tensorflow_probability/python/bijectors/cumsum.py

@@ -20,16 +20,14 @@
 
 import tensorflow.compat.v2 as tf
 from tensorflow_probability.python.bijectors import bijector
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import prefer_static
 
 __all__ = [
     'Cumsum',
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Cumsum(bijector.AutoCompositeTensorBijector):
   """Computes the cumulative sum of a tensor along a specified axis.
 

tensorflow_probability/python/bijectors/discrete_cosine_transform.py

@@ -21,7 +21,6 @@
 import tensorflow.compat.v2 as tf
 
 from tensorflow_probability.python.bijectors import bijector
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 
 
@@ -30,8 +29,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class DiscreteCosineTransform(bijector.AutoCompositeTensorBijector):
   """Compute `Y = g(X) = DCT(X)`, where DCT type is indicated by the `type` arg.