Add an option to convert_to_nested_tensor to not convert objects with reference semantics.

PiperOrigin-RevId: 476461510

Author: emilyfertig

tensorflow_probability/python/internal/BUILD

@@ -420,6 +420,7 @@ multi_substrate_py_library(
     srcs = ["nest_util.py"],
     deps = [
         ":prefer_static",
+        ":tensor_util",
         # numpy dep,
         # tensorflow dep,
     ],

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

@@ -1813,6 +1813,11 @@ def test_convert_variable_to_tensor(self):
     self.assertEqual(np.float64, x.dtype)
     self.assertAllEqual([0., 1., 2.], x)
 
+  def test_convert_variable_to_tensor_with_dtype_hint(self):
+    v = nptf.Variable(np.int32(0))
+    x = nptf.convert_to_tensor(v, dtype_hint=tf.float32)
+    self.assertEqual(np.int32, x.dtype)
+
   def test_get_static_value(self):
     x = nptf.get_static_value(nptf.zeros((3, 2), dtype=nptf.float32))
     self.assertEqual(onp.ndarray, type(x))

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

@@ -141,7 +141,9 @@ def _convert_to_tensor(value, dtype=None, dtype_hint=None, name=None):  # pylint
   """Emulates tf.convert_to_tensor."""
   dtype = utils.numpy_dtype(dtype)
   dtype_hint = utils.numpy_dtype(dtype_hint)
-  if is_tensor(value) and not isinstance(value, Variable):
+  if isinstance(value, Variable):
+    value = value.__wrapped__
+  if is_tensor(value):
     # In NumPy mode, we are lenient on the dtype compatibility check because
     # some codepaths rely on flexible conversion from int/float64 to 32.
     if dtype is not None and value.dtype != dtype:
@@ -619,13 +621,18 @@ def __init__(
       v = v.astype(utils.numpy_dtype(dtype))
     super(NumpyVariable, self).__init__(v)
     self._self_name = name
+    self._self_trainable = trainable
     self.initializer = None
   # pylint: enable=unused-argument
 
   @property
   def name(self):
     return self._self_name if self._self_name is not None else str(id(self))
 
+  @property
+  def trainable(self):
+    return self._self_trainable
+
   def __array__(self, dtype=None):
     if dtype is not None:
       dtype = utils.numpy_dtype(dtype)
@@ -658,10 +665,6 @@ def assign_sub(self, value, **_):
       jax.core.pytype_aval_mappings[onp.ndarray])
 
 
-def _convert_variable_to_tensor(value, dtype=None):
-  return convert_to_tensor(value.__wrapped__, dtype=dtype)
-register_tensor_conversion_function(NumpyVariable, _convert_variable_to_tensor)
-
 Variable = NumpyVariable
 
 

tensorflow_probability/python/internal/nest_util.py

@@ -20,6 +20,7 @@
 import tensorflow.compat.v2 as tf
 
 from tensorflow_probability.python.internal import prefer_static as ps
+from tensorflow_probability.python.internal import tensor_util
 
 from tensorflow.python.util import nest  # pylint: disable=g-direct-tensorflow-import
 
@@ -410,7 +411,7 @@ def call_fn(fn, args):
 
 def convert_to_nested_tensor(value, dtype=None, dtype_hint=None,
                              allow_packing=False, as_shape_tensor=False,
-                             name=None):
+                             convert_ref=True, name=None):
   """Converts the given `value` to a (structure of) `Tensor`.
 
   This function converts Python objects of various types to a (structure of)
@@ -432,6 +433,8 @@ def convert_to_nested_tensor(value, dtype=None, dtype_hint=None,
     as_shape_tensor: Optional boolean when if `True` uses
       `prefer_static.convert_to_shape_tensor` instead of `tf.convert_to_tensor`
       for JAX compatibility.
+    convert_ref: Python `bool`, default `True`. If `True`, convert objects with
+      reference semantics to Tensor.
     name: Optional name to use if a new `Tensor` is created. If inputs are
       structured, elements are named accoring to '{name}/{path}.{to}.{elem}'.
 
@@ -467,8 +470,11 @@ def convert_fn(path, value, dtype, dtype_hint, name=None):
       # break the Bijector cache on forward/inverse log det jacobian,
       # because tf.convert_to_tensor is not a no-op thereon.
       return value
-    else:
+    elif convert_ref:
       return tf.convert_to_tensor(value, dtype, dtype_hint, name=name)
+    else:
+      return tensor_util.convert_nonref_to_tensor(
+          value, dtype, dtype_hint, name=name)
 
   ### The following branches only affect naming.
   # For unstructured calls, just use the provided name.

tensorflow_probability/python/internal/nest_util_test.py

@@ -65,6 +65,27 @@ def __repr__(self):
         self.shape, self.dtype, self.name)
 
 
+class VariableSpec(TensorSpec):
+  """Stub for VariableSpec to simplify tests in np/jax backends."""
+
+  def __init__(self, shape, dtype=tf.float32, trainable=True, name=None):
+    super(VariableSpec, self).__init__(shape, dtype, name=name)
+    self.trainable = trainable
+
+  @classmethod
+  def from_value(cls, variable):
+    return cls(
+        variable.shape, variable.dtype, variable.trainable, variable.name)
+
+  def __eq__(self, other):
+    return (super(VariableSpec, self).__eq__(other)
+            and self.trainable == other.trainable)
+
+  def __repr__(self):
+    return 'VariableSpec(shape={}, dtype={}, trainable={}, name={})'.format(
+        self.shape, self.dtype, self.trainable, self.name)
+
+
 class LeafList(list):
   _tfp_nest_expansion_force_leaf = ()
 
@@ -297,12 +318,13 @@ def fn(arg1, arg2):
                    'b': TensorSpec([], tf.int64, name='c2t/b')}
   },{
       'testcase_name': '_tensor_with_hint',
-      'value': [TensorSpec([], tf.int32)],
+      'value': [TensorSpec([], tf.int32, name='tensor')],
       'dtype_hint': [tf.float32],
       'expected': [TensorSpec([], tf.int32, name='tensor')]
   },{
       'testcase_name': '_tensor_struct',
-      'value': [TensorSpec([], tf.int32), TensorSpec([], tf.float32)],
+      'value': [TensorSpec([], tf.int32, name='tensor'),
+                TensorSpec([], tf.float32, name='tensor')],
       'dtype_hint': [tf.float32, tf.float32],
       'expected': [TensorSpec([], tf.int32, name='tensor'),
                    TensorSpec([], tf.float32, name='tensor_1')]
@@ -318,20 +340,46 @@ def fn(arg1, arg2):
       'name': None,
       'expected': [TensorSpec([], tf.float32, name='Const'),
                    TensorSpec([], tf.float32, name='Const_1')]
+  },{
+      'testcase_name': '_tensor_and_variable_struct',
+      'value': [TensorSpec([], tf.int32, name='tensor'),
+                VariableSpec([], tf.float32, trainable=False, name='variable')],
+      'dtype_hint': [tf.float32, tf.float32],
+      'convert_ref': False,
+      'expected': [
+          TensorSpec([], tf.int32, name='tensor'),
+          VariableSpec([], tf.float32, trainable=False, name='variable:0')]
+  },{
+      'testcase_name': '_tensor_and_variable_struct_convert_ref',
+      'value': [VariableSpec([], tf.int32, name='variable'),
+                TensorSpec([], tf.float32, name='tensor')],
+      'dtype_hint': [tf.float32, tf.float32],
+      'expected': [TensorSpec([], tf.int32, name='c2t/ReadVariableOp'),
+                   TensorSpec([], tf.float32, name='tensor')]
   })
   def testConvertToNestedTensor(
-      self, value, dtype=None, dtype_hint=None, name='c2t', expected=None):
-    # Convert specs to tensors
+      self, value, dtype=None, dtype_hint=None, name='c2t', convert_ref=True,
+      expected=None):
+    # Convert specs to tensors or variables.
     def maybe_spec_to_tensor(x):
+      if isinstance(x, VariableSpec):
+        return tf.Variable(
+            tf.zeros(x.shape, x.dtype), trainable=x.trainable, name=x.name)
       if isinstance(x, TensorSpec):
-        return tf.zeros(x.shape, x.dtype, name='tensor')
+        return tf.zeros(x.shape, x.dtype, name=x.name)
       return x
     value = nest.map_structure(maybe_spec_to_tensor, value)
 
     # Grab shape/dtype from convert_to_nested_tensor for comparison.
+    def spec_from_value(x):
+      if isinstance(x, tf.Variable):
+        return VariableSpec.from_value(x)
+      return TensorSpec.from_tensor(x)
+
     observed = nest.map_structure(
-        TensorSpec.from_tensor,
-        nest_util.convert_to_nested_tensor(value, dtype, dtype_hint, name=name))
+        spec_from_value,
+        nest_util.convert_to_nested_tensor(
+            value, dtype, dtype_hint, convert_ref=convert_ref, name=name))
     self.assertAllEqualNested(observed, expected)
 
   @parameterized.named_parameters({