tfp.util.TransformedVariable and tfp.util.DeferredTensor subclass CompositeTensor.

PiperOrigin-RevId: 374959903

Author: emilyfertig

tensorflow_probability/python/util/deferred_tensor.py

@@ -17,6 +17,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import abc
 import functools
 import numpy as np
 import six
@@ -81,10 +82,10 @@ def _tensorize(d, dtype=None, name=None, as_ref=False):
   return d._value(dtype, name, as_ref)  # pylint: disable=protected-access
 
 
-class TensorMetaClass(type):
+class TensorMetaClass(abc.ABCMeta):
   """A type of class which will make objects which act like Tensors."""
 
-  def __new__(mcs, name, bases, attrs):
+  def __new__(mcs, name, bases, attrs):  # pylint: disable=bad-mcs-classmethod-argument
     operators = set(tf.Tensor.OVERLOADABLE_OPERATORS)
     operators.difference_update({'__eq__', '__ne__'})
     operators.update({'__iter__'})
@@ -109,15 +110,16 @@ def __new__(mcs, name, bases, attrs):
       attrs.update(
           (attr, getattr(tf.Tensor, attr))
           for attr in {'__bool__', '__array_priority__', '__nonzero__'})
-    cls = super(TensorMetaClass, mcs).__new__(mcs, name, bases, attrs)
+    cls = super(TensorMetaClass, mcs).__new__(mcs, name, bases, attrs)  # pylint: disable=too-many-function-args
     tf.register_tensor_conversion_function(cls, conversion_func=_tensorize)
     return cls
 
 
 NONE_SPECIFIED = 'None'
 
 
-class DeferredTensor(six.with_metaclass(TensorMetaClass, tf.Module)):
+class DeferredTensor(six.with_metaclass(
+    TensorMetaClass, tf.Module, tf.__internal__.CompositeTensor)):
   """Variable tracking object which applies function upon `convert_to_tensor`.
 
   #### Example
@@ -379,6 +381,39 @@ def __array__(self, dtype=None):
           'numpy array.')
     return np.array(self._value(dtype=dtype))
 
+  def _get_input_spec(self):
+    if isinstance(self.pretransformed_input, tf.__internal__.CompositeTensor):
+      return self.pretransformed_input._type_spec  # pylint: disable=protected-access
+    if isinstance(self.pretransformed_input, tf.Variable):
+      return resource_variable_ops.VariableSpec(
+          self.pretransformed_input.shape,
+          dtype=self.pretransformed_input.dtype,
+          trainable=self.pretransformed_input.trainable)
+    return tf.TensorSpec.from_tensor(self.pretransformed_input)
+
+  @property
+  def _type_spec(self):
+    input_spec = self._get_input_spec()
+    transform_or_spec = getattr(self._transform_fn, '_type_spec',
+                                self._transform_fn)
+
+    # Extract Variables from also_track.
+    if self.also_track is None:
+      also_track_spec = None
+    else:
+      also_track_vars = tf.nest.flatten(
+          tf.nest.map_structure(
+              lambda x: x.variables if isinstance(x, tf.Module) else x,
+              self.also_track))
+      also_track_spec = tf.nest.map_structure(
+          lambda x: resource_variable_ops.VariableSpec(  # pylint: disable=g-long-lambda
+              x.shape, x.dtype, trainable=x.trainable),
+          also_track_vars)
+
+    return _DeferredTensorSpec(
+        input_spec, transform_or_spec, dtype=self.dtype, shape=self.shape,
+        name=self.name, also_track_spec=also_track_spec)
+
 
 class TransformedVariable(DeferredTensor):
   """Variable tracking object which applies a bijector upon `convert_to_tensor`.
@@ -455,11 +490,7 @@ def __init__(self, initial_value, bijector, dtype=None, name=None, **kwargs):
         which is the initial value for the `TransformedVariable`. The underlying
         untransformed `tf.Variable` will be initialized with
         `bijector.inverse(initial_value)`. Can also be a callable with no
-        argument that returns the initial value when called. Note: if
-        `initial_value` is a `TransformedVariable` then the instantiated object
-        does not create a new `tf.Variable`, but rather points to the underlying
-        `Variable` and chains the `bijector` arg with the underlying bijector as
-        `tfb.Chain([bijector, initial_value.bijector])`.
+        argument that returns the initial value when called.
       bijector: A `Bijector`-like instance which defines the transformations
         applied to the underlying `tf.Variable`.
       dtype: `tf.dtype.DType` instance or otherwise valid `dtype` value to
@@ -479,16 +510,25 @@ def __init__(self, initial_value, bijector, dtype=None, name=None, **kwargs):
 
     if callable(initial_value):
       initial_value = initial_value()
-    initial_value = tf.convert_to_tensor(
-        initial_value, dtype_hint=bijector.dtype, dtype=dtype)
+
+    # Extra kwarg that TypeSpec._from_components uses to re-build the object
+    # without re-initializing the variable.
+    pretransformed_input = kwargs.pop('pretransformed_input', None)
+    if pretransformed_input is None:
+      initial_value = tf.convert_to_tensor(
+          initial_value, dtype_hint=bijector.dtype, dtype=dtype)
+      pretransformed_input = tf.Variable(
+          initial_value=bijector.inverse(initial_value),
+          name=name,
+          dtype=dtype,
+          **kwargs)
+      shape = initial_value.shape
+    else:
+      shape = bijector.forward_event_shape(pretransformed_input.shape)
     super(TransformedVariable, self).__init__(
-        pretransformed_input=tf.Variable(
-            initial_value=bijector.inverse(initial_value),
-            name=name,
-            dtype=dtype,
-            **kwargs),
+        pretransformed_input=pretransformed_input,
         transform_fn=bijector,
-        shape=initial_value.shape,
+        shape=shape,
         name=bijector.name)
     self._bijector = bijector
 
@@ -529,6 +569,13 @@ def assign_sub(self, value, use_locking=False, name=None, read_value=True):
         name=name,
         read_value=read_value)
 
+  @property
+  def _type_spec(self):
+    input_spec = self._get_input_spec()
+    transform_or_spec = getattr(self.bijector, '_type_spec', self.bijector)
+    return _TransformedVariableSpec(
+        input_spec, transform_or_spec, self.dtype, self.name)
+
 
 class _DeferredTensorSpecBase(object):
   """Common methods for '_DeferredTensorSpec' and '_TransformedVariableSpec."""

tensorflow_probability/python/util/deferred_tensor_test.py

@@ -35,6 +35,24 @@
 JAX_MODE = False
 
 
+class NonCompositeTensorExp(object):
+
+  name = 'non_composite_exp'
+  dtype = tf.float32
+
+  def forward(self, x):
+    return tf.math.exp(x)
+
+  def inverse(self, y):
+    return tf.math.log(y)
+
+  def forward_event_shape(self, shape):
+    return shape
+
+  def inverse_event_shape(self, shape):
+    return shape
+
+
 @test_util.test_all_tf_execution_regimes
 class DeferredTensorTest(test_util.TestCase):
 
@@ -149,6 +167,55 @@ def test_from_bijector_with_inverted_assignment(self):
       self.assertAllEqual([[-1.], [-2.], [-3.]], v_)
       self.assertAllEqual([[-1., 0.], [-2., 0.], [-3., 0.]], y_)
 
+  @test_util.disable_test_for_backend(
+      disable_numpy=True, disable_jax=True,
+      reason='JAX and Numpy do not have `CompositeTensor`.')
+  @parameterized.named_parameters(
+      ('transform_fn_is_bijector', tfb.Exp),
+      ('transform_fn_is_bijector_like', NonCompositeTensorExp),
+      ('transform_fn_is_callable', lambda: tf.math.exp))
+  def test_composite_tensor(self, make_transform_fn):
+    initial_value = [0.2, 3.]
+    pretransformed_input = tf.Variable(initial_value, dtype=tf.float32)
+    x = tfp.util.DeferredTensor(pretransformed_input, make_transform_fn())
+
+    @tf.function
+    def f(x_):
+      self.assertLen(x_.trainable_variables, 1)
+      return x_
+
+    y = f(x)
+    self.evaluate([v.initializer for v in x.trainable_variables])
+    self.assertAllClose(self.evaluate(tf.convert_to_tensor(y)),
+                        self.evaluate(tf.math.exp(initial_value)))
+    self.assertLen(x.trainable_variables, 1)
+    self.assertLen(y.trainable_variables,
+                   1 if tf.config.functions_run_eagerly() else 0)
+
+  @test_util.disable_test_for_backend(
+      disable_numpy=True, disable_jax=True,
+      reason=('vectorized_map not implemented in Numpy; '
+              '`DeferredTensor` is not a valid JAX type.'))
+  def test_vectorized_map(self):
+    pretransformed_input = tf.Variable(tf.ones([5, 3]))
+    x = tfp.util.DeferredTensor(pretransformed_input, tfb.Scale([5]))
+    y = tf.vectorized_map(lambda v: v + 2., x)
+    self.evaluate([v.initializer for v in x.trainable_variables])
+    self.assertAllClose(self.evaluate(y), 5. * pretransformed_input + 2.)
+
+  @test_util.disable_test_for_backend(
+      disable_numpy=True, disable_jax=True,
+      reason='JAX and Numpy have no notion of `CompositeTensor`.')
+  def test_also_track_through_flatten_unflatten(self):
+    pretransformed_input = tf.Variable(3.)
+    also_track = tfd.Normal(tf.Variable(0.), scale=1.)
+    x = tfp.util.DeferredTensor(pretransformed_input,
+                                tfb.Shift(tf.Variable(2.)),
+                                also_track=also_track)
+    flat = tf.nest.flatten(x, expand_composites=True)
+    unflat = tf.nest.pack_sequence_as(x, flat, expand_composites=True)
+    self.assertLen(unflat.trainable_variables, 3)
+
 
 @test_util.test_all_tf_execution_regimes
 class TransformedVariableTest(test_util.TestCase):
@@ -274,6 +341,41 @@ def test_nested_transformed_variable(self):
     self.assertIsNot(x.pretransformed_input, y.pretransformed_input)
     # Different vars have no deps so we needn't test cross-talk.
 
+  @test_util.disable_test_for_backend(
+      disable_numpy=True, disable_jax=True,
+      reason='JAX and Numpy do not have `CompositeTensor`.')
+  @parameterized.named_parameters(
+      ('composite_bijector', tfb.Softplus),
+      ('non_composite_bijector', NonCompositeTensorExp))
+  def test_composite_tensor(self, make_bijector):
+    x = tfp.util.TransformedVariable(5., make_bijector())
+    add_val = 10.
+
+    @tf.function
+    def f(x_):
+      x_.assign_add(add_val)
+      self.assertLen(x_.trainable_variables, 1)
+      return x_
+
+    y = f(x)
+    self.evaluate([v.initializer for v in x.trainable_variables])
+    self.assertAllClose(self.evaluate(tf.convert_to_tensor(y)), 15.)
+    self.assertAllClose(self.evaluate(tf.convert_to_tensor(x)), 15.)
+    self.assertLen(x.trainable_variables, 1)
+    self.assertLen(y.trainable_variables,
+                   1 if tf.config.functions_run_eagerly() else 0)
+
+  @test_util.disable_test_for_backend(
+      disable_numpy=True, disable_jax=True,
+      reason=('vectorized_map not implemented in Numpy; '
+              '`DeferredTensor` is not a valid JAX type.'))
+  def test_vectorized_map(self):
+    initial_value = tf.ones([5, 3])
+    x = tfp.util.TransformedVariable(initial_value, tfb.Sigmoid())
+    y = tf.vectorized_map(lambda v: v + 2., x)
+    self.evaluate([v.initializer for v in x.trainable_variables])
+    self.assertAllClose(self.evaluate(y), initial_value + 2.)
+
 
 @test_util.test_all_tf_execution_regimes
 class DeferredTensorBehavesLikeTensorTest(test_util.TestCase):