Fix TensorFlow checkpoint and trackable imports.

PiperOrigin-RevId: 457415636

Author: tf-transform-team

tensorflow_transform/annotators.py

@@ -27,8 +27,13 @@
 # pylint: disable=g-direct-tensorflow-import
 from tensorflow.python.framework import func_graph
 from tensorflow.python.framework import ops
-from tensorflow.python.training.tracking import base
-# pylint: enable=g-direct-tensorflow-import
+# pylint: disable=g-import-not-at-top
+try:
+  # Moved in TensorFlow 2.10.
+  from tensorflow.python.trackable import base
+except ImportError:
+  from tensorflow.python.training.tracking import base
+# pylint: enable=g-direct-tensorflow-import, g-import-not-at-top
 
 __all__ = ['annotate_asset', 'make_and_track_object']
 

tensorflow_transform/annotators_test.py

@@ -17,7 +17,13 @@
 from tensorflow_transform import annotators
 from tensorflow_transform import test_case
 
-from tensorflow.python.training.tracking import base  # pylint: disable=g-direct-tensorflow-import
+# pylint: disable=g-direct-tensorflow-import, g-import-not-at-top
+try:
+  # Moved in TensorFlow 2.10.
+  from tensorflow.python.trackable import base
+except ImportError:
+  from tensorflow.python.training.tracking import base
+# pylint: enable=g-direct-tensorflow-import, g-import-not-at-top
 
 
 class AnnotatorsTest(test_case.TransformTestCase):

tensorflow_transform/saved/saved_transform_io_v2.py

@@ -30,9 +30,16 @@
 from tensorflow.python.eager import wrap_function
 from tensorflow.python.framework import composite_tensor
 from tensorflow.python.ops import lookup_ops
-from tensorflow.python.training.tracking import tracking
 from tensorflow.python.util import object_identity
-# pylint: enable=g-direct-tensorflow-import
+
+# pylint: disable=g-import-not-at-top
+try:
+  # Moved in TensorFlow 2.10.
+  from tensorflow.python.trackable import resource as tracking
+except ImportError:
+  from tensorflow.python.training.tracking import tracking
+
+# pylint: enable=g-direct-tensorflow-import, g-import-not-at-top
 
 
 def _restore_from_v1_saved_model(