Add _convert_variables_to_tensors method to TFP library CompositeTensors.

PiperOrigin-RevId: 454701997

Author: emilyfertig

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/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')

tensorflow_probability/python/distributions/platform_compatibility_test.py

@@ -222,6 +222,10 @@
     'WishartTriL': 1e-5,
 })
 
+COMPOSITE_TENSOR_LOGPROB_RTOL = collections.defaultdict(lambda: 1e-6)
+COMPOSITE_TENSOR_LOGPROB_RTOL.update({
+    'WishartTriL': 1e-5,
+})
 
 SKIP_KL_CHECK_DIST_VAR_GRADS = [
     'GeneralizedExtremeValue',  # TD's KL gradients do not rely on bijector
@@ -419,29 +423,35 @@ def _test_sample_and_log_prob(self, dist_name, dist):
     ct_lp1 = unflat.log_prob(sample1)
     orig_lp1 = dist.log_prob(sample1)
     ct_lp1_, orig_lp1_ = self.evaluate((ct_lp1, orig_lp1))
-    self.assertAllClose(ct_lp1_, orig_lp1_)
+    self.assertAllClose(ct_lp1_, orig_lp1_,
+                        rtol=COMPOSITE_TENSOR_LOGPROB_RTOL[dist_name])
 
     # ... and after.  (Even though they're supposed to be the same anyway.)
     ct_lp2 = unflat.log_prob(sample2)
     orig_lp2 = dist.log_prob(sample2)
     ct_lp2_, orig_lp2_ = self.evaluate((ct_lp2, orig_lp2))
-    self.assertAllClose(ct_lp2_, orig_lp2_)
+    self.assertAllClose(ct_lp2_, orig_lp2_,
+                        rtol=COMPOSITE_TENSOR_LOGPROB_RTOL[dist_name])
 
-  # TODO(alexeev): Add coverage for meta distributions, in addition to base
-  # distributions.
   @parameterized.named_parameters(
-      {'testcase_name': dname, 'dist_name': dname}
-      for dname in sorted(list(dhps.INSTANTIABLE_BASE_DISTS.keys()))
+      {'testcase_name': dname, 'dist_name': dname}  # pylint: disable=g-complex-comprehension
+      for dname in sorted(list(dhps.INSTANTIABLE_BASE_DISTS.keys())
+                          + list(dhps.INSTANTIABLE_META_DISTS))
       if dname not in dhps.TF2_UNFRIENDLY_DISTS)
   @hp.given(hps.data())
   @tfp_hps.tfp_hp_settings()
   def testCompositeTensor(self, dist_name, data):
     dist = data.draw(
         dhps.distributions(
-            dist_name=dist_name, enable_vars=False, validate_args=False))
+            dist_name=dist_name, enable_vars=True, validate_args=False,
+            eligibility_filter=(
+                lambda d: type(d).__name__ not in dhps.TF2_UNFRIENDLY_DISTS)))
+    self.evaluate([v.initializer for v in dist.trainable_variables])
     with tfp_hps.no_tf_rank_errors():
       self._test_sample_and_log_prob(dist_name, dist)
 
+    self.assertConvertVariablesToTensorsWorks(dist)
+
 
 @test_util.test_graph_mode_only
 class DistributionXLATest(test_util.TestCase):

tensorflow_probability/python/internal/BUILD

@@ -707,6 +707,7 @@ multi_substrate_py_library(
         "//tensorflow_probability/python/bijectors:bijector",
         "//tensorflow_probability/python/internal:empirical_statistical_testing",
         "//tensorflow_probability/python/internal/backend/numpy",
+        "//tensorflow_probability/python/util:deferred_tensor",
         "//tensorflow_probability/python/util:seed_stream",
     ],
 )

tensorflow_probability/python/internal/auto_composite_tensor.py

@@ -26,6 +26,7 @@
 from tensorflow.python.framework import type_spec
 from tensorflow.python.ops import resource_variable_ops
 from tensorflow.python.util import tf_inspect
+from tensorflow.python.util import variable_utils
 # pylint: enable=g-direct-tensorflow-import
 
 __all__ = [
@@ -527,6 +528,9 @@ def auto_composite_tensor(
     - object.attribute = [tf.constant(1.), [tf.constant(2.)]]   # valid
     - object.attribute = ['abc', tf.constant(1.)]               # invalid
 
+  All `__init__` args that may be `ResourceVariable`s must also admit `Tensor`s
+  (or else `_convert_variables_to_tensors` must be overridden).
+
   If the attribute is a callable, serialization of the `TypeSpec`, and therefore
   interoperability with `tf.saved_model`, is not currently supported. As a
   workaround, callables that do not contain or close over `Tensor`s may be
@@ -655,6 +659,37 @@ def body(obj):
   type_spec_class_name = f'{cls.__name__}_ACTTypeSpec'
   type_spec_name = f'{module_name}.{type_spec_class_name}'
 
+  def _convert_variables_to_tensors(obj):
+    """Recursively converts Variables in the AutoCompositeTensor to Tensors.
+
+    This method flattens `obj` into a nested structure of `Tensor`s or
+    `CompositeTensor`s, converts any `ResourceVariable`s (which are
+    `CompositeTensor`s) to `Tensor`s, and rebuilds `obj` with `Tensor`s in place
+    of `ResourceVariable`s.
+
+    The usage of `obj._type_spec._from_components` violates the contract of
+    `CompositeTensor`, since it is called on a different nested structure
+    (one containing only `Tensor`s) than `obj.type_spec` specifies (one that may
+    contain `ResourceVariable`s). Since `AutoCompositeTensor`'s
+    `_from_components` method passes the contents of the nested structure to
+    `__init__` to rebuild the TFP object, and any TFP object that may be
+    instantiated with `ResourceVariables` may also be instantiated with
+    `Tensor`s, this usage is valid.
+
+    Args:
+      obj: An `AutoCompositeTensor` instance.
+
+    Returns:
+      tensor_obj: `obj` with all internal `ResourceVariable`s converted to
+        `Tensor`s.
+    """
+    # pylint: disable=protected-access
+    components = obj._type_spec._to_components(obj)
+    tensor_components = variable_utils.convert_variables_to_tensors(
+        components)
+    return obj._type_spec._from_components(tensor_components)
+    # pylint: enable=protected-access
+
   # If the declared class is already a CompositeTensor subclass, we can avoid
   # affecting the actual type of the returned class. Otherwise, we need to
   # explicitly mix in the CT type, and hence create and return a newly
@@ -674,7 +709,10 @@ def _type_spec(obj):
       return _AlreadyCTTypeSpec.from_instance(
           obj, omit_kwargs, non_identifying_kwargs)
 
-    cls._type_spec = property(_type_spec)  # pylint: disable=protected-access
+    # pylint: disable=protected-access
+    cls._type_spec = property(_type_spec)
+    cls._convert_variables_to_tensors = _convert_variables_to_tensors
+    # pylint: enable=protected-access
     return cls
 
   clsid = (cls.__module__, cls.__name__, omit_kwargs,
@@ -701,6 +739,9 @@ def _type_spec(self):
       return _GeneratedCTTypeSpec.from_instance(
           self, omit_kwargs, non_identifying_kwargs)
 
+    def _convert_variables_to_tensors(self):
+      return _convert_variables_to_tensors(self)
+
   _AutoCompositeTensor.__name__ = cls.__name__
   _registry[clsid] = _AutoCompositeTensor
   type_spec_register(type_spec_name)(_GeneratedCTTypeSpec)

tensorflow_probability/python/internal/auto_composite_tensor_test.py

@@ -455,6 +455,40 @@ def fn(x):
     batch_dist = tf.vectorized_map(fn, tf.convert_to_tensor([1., 2., 3.]))
     self.assertAllEqual(batch_dist.batch_shape, [3, 2])
 
+  def test_convert_variables_to_tensors(self):
+    v = tf.Variable(0.)
+    u = tf.Variable(1.)
+    var_td = tfd.TransformedDistribution(
+        tfd.Normal(v, 1.),
+        bijector=tfb.Shift(u))
+    tensor_td = var_td._convert_variables_to_tensors()
+
+    self.evaluate([x.initializer for x in var_td.trainable_variables])
+    self.assertIsInstance(var_td, tf.__internal__.CompositeTensor)
+    self.assertLen(var_td.trainable_variables, 2)
+    self.assertEmpty(tensor_td.trainable_variables)
+    self.assertEqual(self.evaluate(var_td.distribution.loc),
+                     self.evaluate(tensor_td.distribution.loc))
+    self.assertEqual(self.evaluate(var_td.bijector.shift),
+                     self.evaluate(tensor_td.bijector.shift))
+
+  def test_automatic_conversion_to_tensor(self):
+    v = tf.Variable(tf.ones([5]))
+    d = tfd.Normal(tf.zeros([5]), v)
+    x = tf.convert_to_tensor([3.])
+
+    vectorized_log_prob = tf.vectorized_map(lambda z: z.log_prob(x), d)
+    log_prob = d.log_prob(x)
+    self.evaluate(v.initializer)
+    self.assertAllClose(vectorized_log_prob[:, 0], log_prob)
+
+    loc = tf.Variable(0.)
+    self.evaluate(loc.initializer)
+    cond_dist = tf.cond(
+        tf.convert_to_tensor(True),
+        lambda: tfd.Normal(loc, 1.), lambda: tfd.Normal(0., 1.))
+    self.assertIsInstance(cond_dist, tfd.Normal)
+
 
 class _TestTypeSpec(auto_composite_tensor._AutoCompositeTensorTypeSpec):
 

tensorflow_probability/python/internal/backend/jax/BUILD

@@ -63,6 +63,7 @@ FILENAMES = [
     "test_lib",
     "tf_inspect",
     "type_spec",
+    "variable_utils",
     "variables",
     "v1",
     "v2",

tensorflow_probability/python/internal/backend/numpy/BUILD

@@ -62,6 +62,7 @@ py_library(
         ":test_lib",
         ":tf_inspect",
         ":type_spec",
+        ":variable_utils",
         ":variables",
     ],
 )
@@ -408,6 +409,11 @@ py_library(
     srcs = ["type_spec.py"],
 )
 
+py_library(
+    name = "variable_utils",
+    srcs = ["variable_utils.py"],
+)
+
 py_library(
     name = "variables",
     srcs = ["variables.py"],

tensorflow_probability/python/internal/backend/numpy/tensor_spec.py

@@ -19,11 +19,17 @@
 ]
 
 
-class TensorSpec(object):
+class DenseSpec(object):
 
-  def __init__(self, shape, dtype):
+  def __init__(self, shape, dtype, name=None):
     self.shape = shape
     self.dtype = dtype
+    self.name = name
 
   def __repr__(self):
-    return f'TensorSpec(shape={self.shape}, dtype={self.dtype})'
+    return '{}(shape={}, dtype={}, name={})'.format(
+        type(self).__name__, self.shape, repr(self.dtype), repr(self.name))
+
+
+class TensorSpec(DenseSpec):
+  pass

tensorflow_probability/python/internal/backend/numpy/variable_utils.py

@@ -0,0 +1,23 @@
+# Copyright 2022 The TensorFlow Probability Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Numpy stub for `variable_utils`."""
+
+__all__ = [
+    'convert_variables_to_tensors',
+]
+
+
+def convert_variables_to_tensors(x):
+  return x