From fdb76cd1e4eb9d985703c5f376234d9c7392a5b2 Mon Sep 17 00:00:00 2001 From: Fiona Lang Date: Fri, 8 Sep 2023 20:33:47 -0700 Subject: [PATCH] Move tpu import from python/__init__.py to python/modules_with_exports.py. PiperOrigin-RevId: 563919380 --- tensorflow/compiler/jit/ops/BUILD | 4 ---- tensorflow/python/BUILD | 1 - tensorflow/python/__init__.py | 1 + tensorflow/python/modules_with_exports.py | 3 --- tensorflow/python/ops/BUILD | 8 -------- tensorflow/python/ops/gradients_impl.py | 7 ------- tensorflow/python/ops/state_ops.py | 1 - .../saved_model/registration/tf_registration_test.py | 3 --- tensorflow/python/tools/api/generator/doc_srcs.py | 8 ++------ tensorflow/python/tpu/ops/tpu_ops.py | 3 --- tensorflow/python/tpu/tpu.py | 2 ++ 11 files changed, 5 insertions(+), 36 deletions(-) diff --git a/tensorflow/compiler/jit/ops/BUILD b/tensorflow/compiler/jit/ops/BUILD index 08798aa9f30..ab93b126fa3 100644 --- a/tensorflow/compiler/jit/ops/BUILD +++ b/tensorflow/compiler/jit/ops/BUILD @@ -36,9 +36,5 @@ py_strict_library( name = "xla_ops_grad", srcs = ["xla_ops_grad.py"], srcs_version = "PY3", - visibility = [ - "//tensorflow/compiler/tf2xla:internal", - "//tensorflow/python/ops:__pkg__", - ], deps = ["//tensorflow/python/framework:ops"], ) diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD index cf190be79f9..55c421f5202 100644 --- a/tensorflow/python/BUILD +++ b/tensorflow/python/BUILD @@ -438,7 +438,6 @@ py_strict_library( "//tensorflow/python/profiler:trace", "//tensorflow/python/saved_model", "//tensorflow/python/summary:summary_py", - "//tensorflow/python/tpu:tpu_noestimator", "//tensorflow/python/training", "//tensorflow/python/training:quantize_training", "//tensorflow/python/user_ops:ops", diff --git a/tensorflow/python/__init__.py b/tensorflow/python/__init__.py index 55fe1cbdcb7..cc24e60c698 100644 --- a/tensorflow/python/__init__.py +++ b/tensorflow/python/__init__.py @@ -34,6 +34,7 @@ import traceback # from tensorflow.python import keras # from tensorflow.python.layers import layers +from tensorflow.python.tpu import api # Special dunders that we choose to export: _exported_dunders = set([ diff --git a/tensorflow/python/modules_with_exports.py b/tensorflow/python/modules_with_exports.py index 03acc30f4ba..6e59fc7f944 100644 --- a/tensorflow/python/modules_with_exports.py +++ b/tensorflow/python/modules_with_exports.py @@ -169,9 +169,6 @@ from tensorflow.python.client.client_lib import * # Summary from tensorflow.python.summary import summary -# TPU -from tensorflow.python.tpu import api - # Training from tensorflow.python.training import training as train from tensorflow.python.training import quantize_training as _quantize_training diff --git a/tensorflow/python/ops/BUILD b/tensorflow/python/ops/BUILD index 05fdc166033..3fdf838a332 100644 --- a/tensorflow/python/ops/BUILD +++ b/tensorflow/python/ops/BUILD @@ -1620,7 +1620,6 @@ py_strict_library( ":array_ops", ":check_ops", ":control_flow_grad", - ":cudnn_rnn_grad", ":gradients_util", ":image_grad", ":linalg_grad", @@ -1629,20 +1628,14 @@ py_strict_library( ":manip_grad", ":math_grad", ":math_ops", - ":nccl_ops", ":nn_grad", ":optional_grad", - ":proto_ops", ":random_grad", ":rnn_grad", ":sdca_ops", - ":sets", - ":sparse_grad", - ":tensor_array_grad", ":tensor_array_ops", ":unconnected_gradients", ":while_loop", - "//tensorflow/compiler/jit/ops:xla_ops_grad", "//tensorflow/python/debug/lib:debug_gradients", "//tensorflow/python/debug/lib:dumping_callback", "//tensorflow/python/framework:dtypes", @@ -2915,7 +2908,6 @@ py_strict_library( ":array_ops", ":math_ops_gen", ":resource_variable_ops_gen", - ":state_grad", ":state_ops_gen", "//tensorflow/python/framework:ops", "//tensorflow/python/framework:tensor_shape", diff --git a/tensorflow/python/ops/gradients_impl.py b/tensorflow/python/ops/gradients_impl.py index ef4b74e07be..b47f8d8bf9e 100644 --- a/tensorflow/python/ops/gradients_impl.py +++ b/tensorflow/python/ops/gradients_impl.py @@ -14,7 +14,6 @@ # ============================================================================== """Implements the graph generation for computation of gradients.""" -from tensorflow.compiler.jit.ops import xla_ops_grad # pylint: disable=unused-import from tensorflow.python.debug.lib import debug_gradients # pylint: disable=unused-import from tensorflow.python.debug.lib import dumping_callback # pylint: disable=unused-import from tensorflow.python.framework import dtypes @@ -23,7 +22,6 @@ from tensorflow.python.ops import array_grad # pylint: disable=unused-import from tensorflow.python.ops import array_ops from tensorflow.python.ops import check_ops # pylint: disable=unused-import from tensorflow.python.ops import control_flow_grad # pylint: disable=unused-import -from tensorflow.python.ops import cudnn_rnn_grad # pylint: disable=unused-import from tensorflow.python.ops import gradients_util from tensorflow.python.ops import image_grad # pylint: disable=unused-import from tensorflow.python.ops import linalg_grad # pylint: disable=unused-import @@ -32,16 +30,11 @@ from tensorflow.python.ops import logging_ops # pylint: disable=unused-import from tensorflow.python.ops import manip_grad # pylint: disable=unused-import from tensorflow.python.ops import math_grad # pylint: disable=unused-import from tensorflow.python.ops import math_ops -from tensorflow.python.ops import nccl_ops # pylint: disable=unused-import from tensorflow.python.ops import nn_grad # pylint: disable=unused-import from tensorflow.python.ops import optional_grad # pylint: disable=unused-import -from tensorflow.python.ops import proto_ops # pylint: disable=unused-import from tensorflow.python.ops import random_grad # pylint: disable=unused-import from tensorflow.python.ops import rnn_grad # pylint: disable=unused-import from tensorflow.python.ops import sdca_ops # pylint: disable=unused-import -from tensorflow.python.ops import sets # pylint: disable=unused-import -from tensorflow.python.ops import sparse_grad # pylint: disable=unused-import -from tensorflow.python.ops import tensor_array_grad # pylint: disable=unused-import from tensorflow.python.ops import tensor_array_ops from tensorflow.python.ops import while_loop from tensorflow.python.ops.linalg.sparse import sparse_csr_matrix_grad # pylint: disable=unused-import diff --git a/tensorflow/python/ops/state_ops.py b/tensorflow/python/ops/state_ops.py index b80487cd78c..d4cfb23145c 100644 --- a/tensorflow/python/ops/state_ops.py +++ b/tensorflow/python/ops/state_ops.py @@ -24,7 +24,6 @@ from tensorflow.python.ops import array_ops from tensorflow.python.ops import gen_math_ops from tensorflow.python.ops import gen_resource_variable_ops from tensorflow.python.ops import gen_state_ops -from tensorflow.python.ops import state_grad # pylint: disable=unused-import # go/tf-wildcard-import # pylint: disable=wildcard-import from tensorflow.python.ops.gen_state_ops import * diff --git a/tensorflow/python/saved_model/registration/tf_registration_test.py b/tensorflow/python/saved_model/registration/tf_registration_test.py index 7a162e8b3f2..68619f089ed 100644 --- a/tensorflow/python/saved_model/registration/tf_registration_test.py +++ b/tensorflow/python/saved_model/registration/tf_registration_test.py @@ -31,9 +31,6 @@ SERIALIZABLE_ALLOWLIST = os.path.join(resource_loader.get_data_files_path(), CHECKPOINT_SAVER_ALLOWLIST = os.path.join(resource_loader.get_data_files_path(), "tf_checkpoint_saver_allowlist.txt") -# call TPUEmbedding to load the file and run its registration -_ = tf.tpu.experimental.embedding.TPUEmbedding - @registration.register_tf_serializable() class NotInAllowlistExample(tf.Variable): diff --git a/tensorflow/python/tools/api/generator/doc_srcs.py b/tensorflow/python/tools/api/generator/doc_srcs.py index 8749788ed1f..43c468e98d6 100644 --- a/tensorflow/python/tools/api/generator/doc_srcs.py +++ b/tensorflow/python/tools/api/generator/doc_srcs.py @@ -61,12 +61,8 @@ _TENSORFLOW_DOC_SOURCES = { DocSource(docstring_module_name='ops.linalg_ops'), 'logging': DocSource(docstring_module_name='ops.logging_ops'), - 'losses': DocSource( - docstring=( - 'Loss operations for use in neural networks. Note: All the losses' - ' are added to the `GraphKeys.LOSSES` collection by default.' - ) - ), + 'losses': + DocSource(docstring_module_name='ops.losses.losses'), 'manip': DocSource(docstring_module_name='ops.manip_ops'), 'math': diff --git a/tensorflow/python/tpu/ops/tpu_ops.py b/tensorflow/python/tpu/ops/tpu_ops.py index d54cb5a68af..629f5269e1a 100644 --- a/tensorflow/python/tpu/ops/tpu_ops.py +++ b/tensorflow/python/tpu/ops/tpu_ops.py @@ -26,9 +26,6 @@ from tensorflow.python.tpu import tpu_function from tensorflow.python.util.tf_export import tf_export -ops.NotDifferentiable("TPUReplicatedInput") - - def _create_default_group_assignment(): num_shards = tpu_function.get_tpu_context().number_of_shards if num_shards is None: diff --git a/tensorflow/python/tpu/tpu.py b/tensorflow/python/tpu/tpu.py index cc830a68dc4..2daee6c2a96 100644 --- a/tensorflow/python/tpu/tpu.py +++ b/tensorflow/python/tpu/tpu.py @@ -59,6 +59,8 @@ from tensorflow.python.util import variable_utils from tensorflow.python.util.tf_export import tf_export +ops.NotDifferentiable("TPUReplicatedInput") + # Ops which can be safely pruned from XLA compile if they have no consumers. # These ops should also have no inputs. _UNCONNECTED_OPS_TO_PRUNE = set(["Placeholder", "VarHandleOp"]) -- GitLab