Fix infinite loop when converting tf1.Dimension to a nptf.Tensor.

It was caused by using nptf.dimension_value on a tf.Dimension, which did not correctly strip away the Dimension type, since nptf.dimension was checking for type equality against nptf.Dimension (rather than tf.Dimension).

PiperOrigin-RevId: 398351764

Author: SiegeLordEx

tensorflow_probability/python/internal/prefer_static.py

@@ -43,11 +43,20 @@
 
 JAX_MODE = False
 
+
 # Enable converting TF TensorShape and Dimension into np.array. This allows TF
 # code to pass TensorShapes into prefer_static functions. We can also re-use the
-# nptf methods.
+# TensorShape conversion logic from nptf, but Dimension needs to be here to make
+# sure we use TF's dimension_value.
+def _convert_dimension_to_tensor(value, dtype=None):
+  dtype = dtype or np.int32
+  if dtype not in (np.int32, np.int64):
+    raise nptf.ops.TypeConversionError(value, dtype)
+  return nptf.convert_to_tensor(tf.compat.dimension_value(value), dtype=dtype)
+
+
 nptf.register_tensor_conversion_function(
-    tf1.Dimension, nptf.ops._convert_dimension_to_tensor)  # pylint: disable=protected-access
+    tf1.Dimension, _convert_dimension_to_tensor)
 nptf.register_tensor_conversion_function(
     tf.TensorShape, nptf.ops._convert_tensorshape_to_tensor)  # pylint: disable=protected-access
 

tensorflow_probability/python/internal/prefer_static_test.py

@@ -267,6 +267,10 @@ def test_rank_from_shape_scalar(self):
     self.evaluate(v.initializer)
     self.assertEqual(1, self.evaluate(ps.rank_from_shape(v)))
 
+  def test_convert_dimension_to_tensor(self):
+    v = ps.constant(tf1.Dimension(1))
+    self.assertEqual(1, v)
+
   def test_rank_from_shape(self):
     shape = [2, 4, 3]
     expected_rank = len(shape)