diff --git a/tensorflow_probability/python/internal/backend/meta/gen_linear_operators.py b/tensorflow_probability/python/internal/backend/meta/gen_linear_operators.py
index 972bf050d6..4502129771 100644
--- 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
index e2c7f5ccb5..30dbac9d28 100644
--- 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.