Fix TensorFlow checkpoint and trackable imports.
PiperOrigin-RevId: 456866415
diff --git a/tensorflow/cc/saved_model/BUILD b/tensorflow/cc/saved_model/BUILD
index abc03d8..671a1a9 100644
--- a/tensorflow/cc/saved_model/BUILD
+++ b/tensorflow/cc/saved_model/BUILD
@@ -247,7 +247,7 @@
"//tensorflow/python/module",
"//tensorflow/python/saved_model",
"//tensorflow/python/saved_model:save_options",
- "//tensorflow/python/training/tracking",
+ "//tensorflow/python/trackable:asset",
"@absl_py//absl:app",
],
)
diff --git a/tensorflow/cc/saved_model/testdata/generate_saved_models.py b/tensorflow/cc/saved_model/testdata/generate_saved_models.py
index baaa1b1..7990de8 100644
--- a/tensorflow/cc/saved_model/testdata/generate_saved_models.py
+++ b/tensorflow/cc/saved_model/testdata/generate_saved_models.py
@@ -31,7 +31,7 @@
from tensorflow.python.platform import test
from tensorflow.python.saved_model import save_options
from tensorflow.python.saved_model import saved_model
-from tensorflow.python.training.tracking import tracking
+from tensorflow.python.trackable import asset
class VarsAndArithmeticObjectGraph(module.Module):
@@ -71,7 +71,7 @@
class AssetModule(module.Module):
def __init__(self):
- self.asset = tracking.Asset(
+ self.asset = asset.Asset(
test.test_src_dir_path("cc/saved_model/testdata/test_asset.txt"))
@def_function.function(input_signature=[])
@@ -83,7 +83,7 @@
"""A module with an Asset, StaticHashTable, and a lookup function."""
def __init__(self):
- self.asset = tracking.Asset(
+ self.asset = asset.Asset(
test.test_src_dir_path(
"cc/saved_model/testdata/static_hashtable_asset.txt"))
self.table = lookup_ops.StaticHashTable(
diff --git a/tensorflow/lite/tools/optimize/debugging/python/BUILD b/tensorflow/lite/tools/optimize/debugging/python/BUILD
index e55042f..8d5a744 100644
--- a/tensorflow/lite/tools/optimize/debugging/python/BUILD
+++ b/tensorflow/lite/tools/optimize/debugging/python/BUILD
@@ -37,7 +37,7 @@
"//tensorflow/lite/python/metrics",
"//tensorflow/python:framework_test_lib",
"//tensorflow/python/platform:client_testlib",
- "//tensorflow/python/training/tracking",
+ "//tensorflow/python/trackable:autotrackable",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
diff --git a/tensorflow/lite/tools/optimize/debugging/python/debugger_test.py b/tensorflow/lite/tools/optimize/debugging/python/debugger_test.py
index 611175e..f66615a 100644
--- a/tensorflow/lite/tools/optimize/debugging/python/debugger_test.py
+++ b/tensorflow/lite/tools/optimize/debugging/python/debugger_test.py
@@ -29,12 +29,12 @@
from tensorflow.lite.tools.optimize.debugging.python import debugger
from tensorflow.python.framework import test_util
from tensorflow.python.platform import test
-from tensorflow.python.training.tracking import tracking
+from tensorflow.python.trackable import autotrackable
def _get_model():
"""Returns somple model with Conv2D and representative dataset gen."""
- root = tracking.AutoTrackable()
+ root = autotrackable.AutoTrackable()
kernel_in = np.array([-2, -1, 1, 2], dtype=np.float32).reshape((2, 2, 1, 1))
@tf.function(
diff --git a/tensorflow/python/checkpoint/BUILD b/tensorflow/python/checkpoint/BUILD
index 9d2a506..b4d3f8b 100644
--- a/tensorflow/python/checkpoint/BUILD
+++ b/tensorflow/python/checkpoint/BUILD
@@ -46,7 +46,6 @@
":util",
"//tensorflow/core:protos_all_py",
"//tensorflow/python:array_ops",
- "//tensorflow/python:checkpoint_management",
"//tensorflow/python:constant_op",
"//tensorflow/python:dtypes",
"//tensorflow/python:errors",
@@ -63,6 +62,7 @@
"//tensorflow/python:util",
"//tensorflow/python:variable_scope",
"//tensorflow/python:variables",
+ "//tensorflow/python/checkpoint:checkpoint_management",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/saved_model:utils",
@@ -85,7 +85,6 @@
":checkpoint",
":checkpoint_options",
":graph_view",
- "//tensorflow/python:checkpoint_management",
"//tensorflow/python:client_testlib",
"//tensorflow/python:constant_op",
"//tensorflow/python:control_flow_ops",
@@ -103,6 +102,7 @@
"//tensorflow/python:training_util",
"//tensorflow/python:variable_scope",
"//tensorflow/python:variables",
+ "//tensorflow/python/checkpoint:checkpoint_management",
"//tensorflow/python/eager:backprop",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
@@ -134,7 +134,7 @@
"//tensorflow/python:variable_scope",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:test",
- "//tensorflow/python/training/tracking",
+ "//tensorflow/python/trackable:autotrackable",
"@six_archive//:six",
],
)
@@ -231,7 +231,6 @@
"//tensorflow/python/training:optimizer",
"//tensorflow/python/training/saving:saveable_object",
"//tensorflow/python/training/saving:saveable_object_util",
- "//tensorflow/python/training/tracking",
],
)
diff --git a/tensorflow/python/checkpoint/benchmarks_test.py b/tensorflow/python/checkpoint/benchmarks_test.py
index bed992d..b1c6eac 100644
--- a/tensorflow/python/checkpoint/benchmarks_test.py
+++ b/tensorflow/python/checkpoint/benchmarks_test.py
@@ -16,7 +16,7 @@
import os
import time
-
+from tensorflow.python.checkpoint import checkpoint as util
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.module import module
@@ -27,7 +27,6 @@
from tensorflow.python.trackable import base
from tensorflow.python.training import py_checkpoint_reader
from tensorflow.python.training.saving import saveable_object
-from tensorflow.python.training.tracking import util
class _TrivialSaveable(saveable_object.SaveableObject):
diff --git a/tensorflow/python/checkpoint/checkpoint_with_v1_optimizers_test.py b/tensorflow/python/checkpoint/checkpoint_with_v1_optimizers_test.py
index 45391e5..cf42474 100644
--- a/tensorflow/python/checkpoint/checkpoint_with_v1_optimizers_test.py
+++ b/tensorflow/python/checkpoint/checkpoint_with_v1_optimizers_test.py
@@ -28,8 +28,8 @@
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import template
from tensorflow.python.ops import variable_scope
+from tensorflow.python.trackable import autotrackable
from tensorflow.python.training import adam
-from tensorflow.python.training.tracking import tracking
class CheckpointingTests(test.TestCase):
@@ -190,7 +190,7 @@
self.assertAllEqual(3., self.evaluate(beta1_power))
-class _ManualScope(tracking.AutoTrackable):
+class _ManualScope(autotrackable.AutoTrackable):
def __call__(self):
with variable_scope.variable_scope("ManualScope") as vs:
diff --git a/tensorflow/python/feature_column/BUILD b/tensorflow/python/feature_column/BUILD
index 10b2dc1..1979c48 100644
--- a/tensorflow/python/feature_column/BUILD
+++ b/tensorflow/python/feature_column/BUILD
@@ -81,8 +81,10 @@
"//tensorflow/python:variable_scope",
"//tensorflow/python:variables",
"//tensorflow/python/eager:context",
- "//tensorflow/python/training/tracking",
- "//tensorflow/python/training/tracking:data_structures",
+ "//tensorflow/python/trackable:autotrackable",
+ "//tensorflow/python/trackable:base",
+ "//tensorflow/python/trackable:data_structures",
+ "//tensorflow/python/trackable:resource",
"//tensorflow/python/util:tf_export",
"//third_party/py/numpy",
"@six_archive//:six",
diff --git a/tensorflow/python/feature_column/feature_column_v2.py b/tensorflow/python/feature_column/feature_column_v2.py
index d5c337a..d10dbfe 100644
--- a/tensorflow/python/feature_column/feature_column_v2.py
+++ b/tensorflow/python/feature_column/feature_column_v2.py
@@ -154,10 +154,10 @@
from tensorflow.python.ops import variables
from tensorflow.python.platform import gfile
from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.trackable import autotrackable
+from tensorflow.python.trackable import base as trackable
+from tensorflow.python.trackable import data_structures
from tensorflow.python.training import checkpoint_utils
-from tensorflow.python.training.tracking import base as trackable
-from tensorflow.python.training.tracking import data_structures
-from tensorflow.python.training.tracking import tracking
from tensorflow.python.util import deprecation
from tensorflow.python.util import nest
from tensorflow.python.util import tf_inspect
@@ -3159,7 +3159,8 @@
'`DenseFeatures` or `LinearModel` instead.')
-class SharedEmbeddingColumnCreator(tracking.AutoTrackable):
+class SharedEmbeddingColumnCreator(autotrackable.AutoTrackable):
+ """Class that creates a `SharedEmbeddingColumn`."""
def __init__(self,
dimension,
diff --git a/tensorflow/python/kernel_tests/nn_ops/BUILD b/tensorflow/python/kernel_tests/nn_ops/BUILD
index e13c612..4109a67 100644
--- a/tensorflow/python/kernel_tests/nn_ops/BUILD
+++ b/tensorflow/python/kernel_tests/nn_ops/BUILD
@@ -575,7 +575,7 @@
"//tensorflow/python/eager:def_function",
"//tensorflow/python/saved_model:load",
"//tensorflow/python/saved_model:save",
- "//tensorflow/python/training/tracking",
+ "//tensorflow/python/trackable:autotrackable",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
diff --git a/tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py b/tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py
index bed3cbf..5a15d7a 100644
--- a/tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py
@@ -21,6 +21,7 @@
import numpy as np
from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.checkpoint import checkpoint as trackable_utils
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.framework import constant_op
@@ -48,8 +49,7 @@
from tensorflow.python.platform import tf_logging
from tensorflow.python.saved_model import load
from tensorflow.python.saved_model import save
-from tensorflow.python.training.tracking import tracking
-from tensorflow.python.training.tracking import util as trackable_utils
+from tensorflow.python.trackable import autotrackable
from tensorflow.python.util import nest
@@ -3094,7 +3094,7 @@
self.skipTest("b/175887901")
with self.cached_session():
- root = tracking.AutoTrackable()
+ root = autotrackable.AutoTrackable()
root.cell = rnn_cell_impl.LSTMCell(8)
@def_function.function(input_signature=[tensor_spec.TensorSpec([3, 8])])
def call(x):
diff --git a/tensorflow/python/module/BUILD b/tensorflow/python/module/BUILD
index 8e41678d..c4ef3e6 100644
--- a/tensorflow/python/module/BUILD
+++ b/tensorflow/python/module/BUILD
@@ -14,7 +14,7 @@
"//tensorflow/python:tf2",
"//tensorflow/python:util",
"//tensorflow/python:variables",
- "//tensorflow/python/training/tracking",
+ "//tensorflow/python/trackable:autotrackable",
"@six_archive//:six",
],
)
diff --git a/tensorflow/python/module/module.py b/tensorflow/python/module/module.py
index 0c7c5a0..36156d4 100644
--- a/tensorflow/python/module/module.py
+++ b/tensorflow/python/module/module.py
@@ -21,14 +21,14 @@
from tensorflow.python import tf2
from tensorflow.python.framework import ops
from tensorflow.python.ops import variables
-from tensorflow.python.training.tracking import tracking
+from tensorflow.python.trackable import autotrackable
from tensorflow.python.util import nest
from tensorflow.python.util import tf_decorator
from tensorflow.python.util.tf_export import tf_export
@tf_export("Module")
-class Module(tracking.AutoTrackable):
+class Module(autotrackable.AutoTrackable):
"""Base neural network module class.
A module is a named container for `tf.Variable`s, other `tf.Module`s and
diff --git a/tensorflow/python/ops/linalg/BUILD b/tensorflow/python/ops/linalg/BUILD
index cd69a88..752becb 100644
--- a/tensorflow/python/ops/linalg/BUILD
+++ b/tensorflow/python/ops/linalg/BUILD
@@ -18,6 +18,7 @@
"//tensorflow/python:tensor_util",
"//tensorflow/python:util",
"//tensorflow/python/ops/signal",
+ "//tensorflow/python/trackable:data_structures",
"//third_party/py/numpy",
],
)
diff --git a/tensorflow/python/ops/linalg/linear_operator.py b/tensorflow/python/ops/linalg/linear_operator.py
index 130eb73..2ac6de4 100644
--- a/tensorflow/python/ops/linalg/linear_operator.py
+++ b/tensorflow/python/ops/linalg/linear_operator.py
@@ -38,7 +38,7 @@
from tensorflow.python.ops.linalg import linear_operator_util
from tensorflow.python.ops.linalg import slicing
from tensorflow.python.platform import tf_logging as logging
-from tensorflow.python.training.tracking import data_structures
+from tensorflow.python.trackable import data_structures
from tensorflow.python.util import deprecation
from tensorflow.python.util import dispatch
from tensorflow.python.util import nest
diff --git a/tensorflow/python/saved_model/registration/BUILD b/tensorflow/python/saved_model/registration/BUILD
index 9bdc095..6aa57f3 100644
--- a/tensorflow/python/saved_model/registration/BUILD
+++ b/tensorflow/python/saved_model/registration/BUILD
@@ -37,7 +37,7 @@
deps = [
":registration",
"//tensorflow/python/eager:test",
- "//tensorflow/python/training/tracking:base",
+ "//tensorflow/python/trackable:base",
"@absl_py//absl/testing:parameterized",
],
)
@@ -54,7 +54,7 @@
"//tensorflow/python:resource_variable_ops",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/eager:test",
- "//tensorflow/python/training/tracking",
+ "//tensorflow/python/trackable:autotrackable",
"@absl_py//absl/testing:parameterized",
],
)
diff --git a/tensorflow/python/saved_model/registration/registration_saving_test.py b/tensorflow/python/saved_model/registration/registration_saving_test.py
index b5fa078..47b24d9 100644
--- a/tensorflow/python/saved_model/registration/registration_saving_test.py
+++ b/tensorflow/python/saved_model/registration/registration_saving_test.py
@@ -20,6 +20,7 @@
from absl.testing import parameterized
from google.protobuf import wrappers_pb2
+from tensorflow.python.checkpoint import checkpoint as util
from tensorflow.python.client import session
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
@@ -35,8 +36,7 @@
from tensorflow.python.saved_model import loader
from tensorflow.python.saved_model import registration
from tensorflow.python.saved_model import save
-from tensorflow.python.training.tracking import tracking
-from tensorflow.python.training.tracking import util
+from tensorflow.python.trackable import autotrackable
@registration.register_serializable()
@@ -57,7 +57,7 @@
@registration.register_serializable()
-class Stack(tracking.AutoTrackable):
+class Stack(autotrackable.AutoTrackable):
def __init__(self, parts=None):
self.parts = parts
@@ -137,7 +137,7 @@
def test_registered_serializable(self, cycles):
@registration.register_serializable(name=f"SaveAndLoad{cycles}")
- class Module(tracking.AutoTrackable):
+ class Module(autotrackable.AutoTrackable):
def __init__(self, name="module"):
self.v = variables.Variable(1.)
@@ -167,7 +167,7 @@
def test_none_proto(self, cycles):
@registration.register_serializable(name=f"NoneProto{cycles}")
- class Module(tracking.AutoTrackable):
+ class Module(autotrackable.AutoTrackable):
def __init__(self, name="module"):
self.v = variables.Variable(1.)
@@ -189,7 +189,7 @@
def test_deserialization_dependencies(self, cycles):
@registration.register_serializable(name=f"Dependency{cycles}")
- class Module(tracking.AutoTrackable):
+ class Module(autotrackable.AutoTrackable):
def __init__(self, v=None):
self.v = v if v is not None else variables.Variable(1.)
@@ -276,7 +276,7 @@
self.assertAllEqual(exported_value, sess.run(value_output.name))
def test_non_strict_predicate(self):
- class NonStrictPredicateClass(tracking.AutoTrackable):
+ class NonStrictPredicateClass(autotrackable.AutoTrackable):
pass
registration.register_checkpoint_saver(
name="NonStrictPredicate",
@@ -289,12 +289,12 @@
ckpt_path = os.path.join(self.get_temp_dir(), "ckpt")
util.Checkpoint(root).write(ckpt_path)
- root2 = tracking.AutoTrackable()
+ root2 = autotrackable.AutoTrackable()
# This should run without throwing an error.
util.Checkpoint(root2).read(ckpt_path)
def test_strict_predicate(self):
- class StrictPredicateClass(tracking.AutoTrackable):
+ class StrictPredicateClass(autotrackable.AutoTrackable):
pass
registration.register_checkpoint_saver(
name="StrictPredicate",
@@ -307,7 +307,7 @@
ckpt_path = os.path.join(self.get_temp_dir(), "ckpt")
util.Checkpoint(root).write(ckpt_path)
- root2 = tracking.AutoTrackable()
+ root2 = autotrackable.AutoTrackable()
with self.assertRaisesRegex(ValueError, "saver cannot be used"):
util.Checkpoint(root2).read(ckpt_path)
@@ -315,7 +315,7 @@
if not context.executing_eagerly():
self.skipTest("This test must run under eager mode.")
- class RestoreClass(tracking.AutoTrackable):
+ class RestoreClass(autotrackable.AutoTrackable):
pass
def save_fn(trackables, file_prefix):
del trackables # Unused.
diff --git a/tensorflow/python/saved_model/registration/registration_test.py b/tensorflow/python/saved_model/registration/registration_test.py
index 1bb7c88..f6bdf5e 100644
--- a/tensorflow/python/saved_model/registration/registration_test.py
+++ b/tensorflow/python/saved_model/registration/registration_test.py
@@ -22,7 +22,7 @@
from tensorflow.python.eager import test
from tensorflow.python.saved_model import registration
-from tensorflow.python.training.tracking import base
+from tensorflow.python.trackable import base
@registration.register_serializable()
diff --git a/tensorflow/python/training/experimental/BUILD b/tensorflow/python/training/experimental/BUILD
index b4a6a4c..fb800ed 100644
--- a/tensorflow/python/training/experimental/BUILD
+++ b/tensorflow/python/training/experimental/BUILD
@@ -21,7 +21,7 @@
"//tensorflow/python/distribute:distribute_lib",
"//tensorflow/python/distribute:reduce_util",
"//tensorflow/python/eager:context",
- "//tensorflow/python/training/tracking:base",
+ "//tensorflow/python/trackable:base",
"//tensorflow/python/util:tf_export",
"@six_archive//:six",
],
@@ -59,12 +59,12 @@
"//tensorflow/python:gradients",
"//tensorflow/python:resource_variable_ops",
"//tensorflow/python:variables",
+ "//tensorflow/python/checkpoint",
"//tensorflow/python/distribute:distribute_lib",
"//tensorflow/python/distribute:mirrored_strategy",
"//tensorflow/python/eager:context",
"//tensorflow/python/training:gradient_descent",
"//tensorflow/python/training:momentum",
- "//tensorflow/python/training/tracking:util",
"@absl_py//absl/testing:parameterized",
],
)
diff --git a/tensorflow/python/training/experimental/loss_scale.py b/tensorflow/python/training/experimental/loss_scale.py
index 0645c05..0c0ef3a 100644
--- a/tensorflow/python/training/experimental/loss_scale.py
+++ b/tensorflow/python/training/experimental/loss_scale.py
@@ -26,7 +26,7 @@
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables
-from tensorflow.python.training.tracking import base as trackable
+from tensorflow.python.trackable import base as trackable
from tensorflow.python.util import deprecation
from tensorflow.python.util import nest
from tensorflow.python.util.tf_export import tf_export
diff --git a/tensorflow/python/training/experimental/loss_scale_optimizer_test.py b/tensorflow/python/training/experimental/loss_scale_optimizer_test.py
index f42df65..d9230fd 100644
--- a/tensorflow/python/training/experimental/loss_scale_optimizer_test.py
+++ b/tensorflow/python/training/experimental/loss_scale_optimizer_test.py
@@ -17,7 +17,7 @@
import os
from absl.testing import parameterized
-
+from tensorflow.python.checkpoint import checkpoint as trackable_utils
from tensorflow.python.distribute import distribution_strategy_context
from tensorflow.python.distribute import mirrored_strategy
from tensorflow.python.eager import context
@@ -33,7 +33,6 @@
from tensorflow.python.training import momentum
from tensorflow.python.training.experimental import loss_scale as loss_scale_module
from tensorflow.python.training.experimental import loss_scale_optimizer
-from tensorflow.python.training.tracking import util as trackable_utils
# If called outside any strategy.scope() calls, this will return the default
# strategy.
