Move tpu import from python/__init__.py to python/modules_with_exports.py.

PiperOrigin-RevId: 563919380

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"],
 )

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

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([

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

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

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

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 *

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):

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

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:

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"])