Revert "Fix TensorFlow checkpoint and trackable imports."

This reverts commit 0a10dd4.
tensorflow · Jun 6, 2022 · bfe7d2b · bfe7d2b
1 parent 6d04325
commit bfe7d2b
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 6 deletions.
diff --git a/tensorflow_probability/python/internal/backend/meta/gen_linear_operators.py b/tensorflow_probability/python/internal/backend/meta/gen_linear_operators.py
@@ -134,11 +134,6 @@ def gen_module(module_name):
       'from tensorflow.python.ops import variables',
       'from tensorflow_probability.python.internal.backend.numpy '
       'import variables')
-  code = code.replace(
-      'from tensorflow.python.trackable '
-      'import data_structures',
-      'from tensorflow_probability.python.internal.backend.numpy '
-      'import data_structures')
   code = code.replace(
       'from tensorflow.python.training.tracking '
       'import data_structures',

diff --git a/tensorflow_probability/python/layers/internal/distribution_tensor_coercible.py b/tensorflow_probability/python/layers/internal/distribution_tensor_coercible.py
@@ -24,7 +24,7 @@
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.util.deferred_tensor import TensorMetaClass
 from tensorflow.python.framework import composite_tensor  # pylint: disable=g-direct-tensorflow-import
-from tensorflow.python.trackable import data_structures  # pylint: disable=g-direct-tensorflow-import
+from tensorflow.python.training.tracking import data_structures  # pylint: disable=g-direct-tensorflow-import
 
 
 __all__ = []  # We intend nothing public.