Fix TensorFlow checkpoint and trackable imports.

PiperOrigin-RevId: 452684705

Author: k-w-w

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',

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.