Fix TensorFlow checkpoint and trackable imports.

PiperOrigin-RevId: 452684705
tensorflow · Jun 3, 2022 · 0a10dd4 · 0a10dd4
1 parent 12e0f6b
commit 0a10dd4
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
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,6 +134,11 @@ 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.training.tracking import data_structures  # pylint: disable=g-direct-tensorflow-import
+from tensorflow.python.trackable import data_structures  # pylint: disable=g-direct-tensorflow-import
 
 
 __all__ = []  # We intend nothing public.