Register _DeferredTensorSpec and _TransformedVariableSpec with auto_composite_tensor.type_spec_register to avoid errors in the type spec registry when reloading TFP in iPython.

PiperOrigin-RevId: 379608557

Author: emilyfertig

tensorflow_probability/python/util/BUILD

@@ -44,6 +44,7 @@ multi_substrate_py_library(
     deps = [
         # numpy dep,
         # tensorflow dep,
+        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:name_util",
         "//tensorflow_probability/python/internal:tensor_util",

tensorflow_probability/python/util/deferred_tensor.py

@@ -23,6 +23,7 @@
 import six
 
 import tensorflow.compat.v2 as tf
+from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import name_util
 from tensorflow_probability.python.internal import tensor_util
@@ -734,7 +735,7 @@ def __hash__(self):
     return hash(self.__get_cmp_key())
 
 
-@type_spec.register('tfp.util.DeferredTensorSpec')
+@auto_composite_tensor.type_spec_register('tfp.util.DeferredTensorSpec')
 class _DeferredTensorSpec(_DeferredTensorSpecBase, type_spec.BatchableTypeSpec):
   """`tf.TypeSpec` for `tfp.util.DeferredTensor`."""
 
@@ -837,7 +838,7 @@ def _unbatch(self):
         also_track_spec=self._also_track_spec)
 
 
-@type_spec.register('tfp.util.TransformedVariableSpec')
+@auto_composite_tensor.type_spec_register('tfp.util.TransformedVariableSpec')
 class _TransformedVariableSpec(
     _DeferredTensorSpecBase, type_spec.BatchableTypeSpec):
   """`tf.TypeSpec` for `tfp.util.TransformedVariable`."""