Migrate references and remove legacy target tpu:tpu_estimator.

PiperOrigin-RevId: 583446480

.github/workflows/build.yml

@@ -8,11 +8,9 @@ jobs:
     steps:
     - uses: actions/checkout@v2
     - name: Set up Python
-      uses: actions/setup-python@v4
+      uses: actions/setup-python@v2
       with:
         python-version: '3.8.x'
-        cache: 'pip'
-        cache-dependency-path: setup.py
     - name: Install dependencies
       run: |
         pip install tf-nightly mock pytest

examples/mnist.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

examples/mnist_dataset.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

examples/toy_model_tpu.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -27,9 +27,9 @@
 from tensorflow.python.data.ops.dataset_ops import Dataset
 from tensorflow.python.platform import flags
 from tensorflow.python.platform import tf_logging as logging
-from tensorflow.python.tpu import tpu_config  # pylint: disable=g-direct-tensorflow-import
-from tensorflow.python.tpu import tpu_estimator  # pylint: disable=g-direct-tensorflow-import
 from tensorflow_estimator.python.estimator import estimator as estimator_lib
+from tensorflow_estimator.python.estimator.tpu import tpu_config  # pylint: disable=g-deprecated-tf-checker
+from tensorflow_estimator.python.estimator.tpu import tpu_estimator  # pylint: disable=g-deprecated-tf-checker
 
 FLAGS = flags.FLAGS
 

mesh_tensorflow/__init__.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/auto_mtf/__init__.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/auto_mtf/api.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/auto_mtf/api_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/auto_mtf/graph_interface.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/auto_mtf/graph_interface_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/auto_mtf/layout_optimizer.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/auto_mtf/layout_optimizer_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/auto_mtf/memory_estimator.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/auto_mtf/memory_estimator_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/auto_mtf/print_cp_model_solution.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/auto_mtf/scheduler.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/auto_mtf/scheduler_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/auto_mtf/valid_layouts.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/auto_mtf/valid_layouts_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/beam_search.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/bert/__init__.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/bert/bert.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/bert/optimization.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/bert/run_classifier.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/bert/run_pretraining.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/bert/run_squad.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/bert/tokenization.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/experimental/__init__.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/experimental/data_aug_lib.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/experimental/data_aug_lib_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/experimental/input_reader.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/experimental/input_reader_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/experimental/model_executor.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/experimental/offline_data_aug.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/experimental/unet.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/import_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/layers.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/layers_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/ops.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/ops_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/ops_with_redefined_builtins.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/optimize.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/placement_mesh_impl.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/simd_mesh_impl.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/simd_mesh_impl_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/test_utils.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/test_utils_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/tpu_variables.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/transformer/__init__.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/transformer/adaptive_softmax.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/transformer/adaptive_softmax_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/transformer/attention.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/transformer/dataset.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mesh_tensorflow/transformer/dataset_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Mesh TensorFlow Authors.
+# Copyright 2023 The Mesh TensorFlow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.