diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 7c814f5e..754f122f 100644
--- a/.github/workflows/build.yml
+++ b/.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
diff --git a/examples/mnist.py b/examples/mnist.py
index 1a6f6f42..96de1530 100644
--- a/examples/mnist.py
+++ b/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.
diff --git a/examples/mnist_dataset.py b/examples/mnist_dataset.py
index c8d7ef26..a51fe2d7 100644
--- a/examples/mnist_dataset.py
+++ b/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.
diff --git a/examples/toy_model_tpu.py b/examples/toy_model_tpu.py
index 80959fdf..bd601e54 100644
--- a/examples/toy_model_tpu.py
+++ b/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
 
diff --git a/mesh_tensorflow/__init__.py b/mesh_tensorflow/__init__.py
index d8eb845d..909a620d 100644
--- a/mesh_tensorflow/__init__.py
+++ b/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.
diff --git a/mesh_tensorflow/auto_mtf/__init__.py b/mesh_tensorflow/auto_mtf/__init__.py
index c5acfbe2..3671fc88 100644
--- a/mesh_tensorflow/auto_mtf/__init__.py
+++ b/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.
diff --git a/mesh_tensorflow/auto_mtf/api.py b/mesh_tensorflow/auto_mtf/api.py
index 37eb0c94..4e7f0623 100644
--- a/mesh_tensorflow/auto_mtf/api.py
+++ b/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.
diff --git a/mesh_tensorflow/auto_mtf/api_test.py b/mesh_tensorflow/auto_mtf/api_test.py
index 199e77d4..0e9ca742 100644
--- a/mesh_tensorflow/auto_mtf/api_test.py
+++ b/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.
diff --git a/mesh_tensorflow/auto_mtf/graph_interface.py b/mesh_tensorflow/auto_mtf/graph_interface.py
index 6170b47e..d9929b31 100644
--- a/mesh_tensorflow/auto_mtf/graph_interface.py
+++ b/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.
diff --git a/mesh_tensorflow/auto_mtf/graph_interface_test.py b/mesh_tensorflow/auto_mtf/graph_interface_test.py
index 974c2a82..0f9ed4e5 100644
--- a/mesh_tensorflow/auto_mtf/graph_interface_test.py
+++ b/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.
diff --git a/mesh_tensorflow/auto_mtf/layout_optimizer.py b/mesh_tensorflow/auto_mtf/layout_optimizer.py
index 420a69ac..09a8bb72 100644
--- a/mesh_tensorflow/auto_mtf/layout_optimizer.py
+++ b/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.
diff --git a/mesh_tensorflow/auto_mtf/layout_optimizer_test.py b/mesh_tensorflow/auto_mtf/layout_optimizer_test.py
index a7af6cef..cacadbe8 100644
--- a/mesh_tensorflow/auto_mtf/layout_optimizer_test.py
+++ b/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.
diff --git a/mesh_tensorflow/auto_mtf/memory_estimator.py b/mesh_tensorflow/auto_mtf/memory_estimator.py
index 0875096a..b8149ded 100644
--- a/mesh_tensorflow/auto_mtf/memory_estimator.py
+++ b/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.
diff --git a/mesh_tensorflow/auto_mtf/memory_estimator_test.py b/mesh_tensorflow/auto_mtf/memory_estimator_test.py
index 14ece7d8..e8ffb033 100644
--- a/mesh_tensorflow/auto_mtf/memory_estimator_test.py
+++ b/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.
diff --git a/mesh_tensorflow/auto_mtf/print_cp_model_solution.py b/mesh_tensorflow/auto_mtf/print_cp_model_solution.py
index 3f10557d..4b45fafa 100644
--- a/mesh_tensorflow/auto_mtf/print_cp_model_solution.py
+++ b/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.
diff --git a/mesh_tensorflow/auto_mtf/scheduler.py b/mesh_tensorflow/auto_mtf/scheduler.py
index eeffc216..0d70178f 100644
--- a/mesh_tensorflow/auto_mtf/scheduler.py
+++ b/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.
diff --git a/mesh_tensorflow/auto_mtf/scheduler_test.py b/mesh_tensorflow/auto_mtf/scheduler_test.py
index 7006278f..c16e8330 100644
--- a/mesh_tensorflow/auto_mtf/scheduler_test.py
+++ b/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.
diff --git a/mesh_tensorflow/auto_mtf/valid_layouts.py b/mesh_tensorflow/auto_mtf/valid_layouts.py
index 9b8d135e..52ad31f4 100644
--- a/mesh_tensorflow/auto_mtf/valid_layouts.py
+++ b/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.
diff --git a/mesh_tensorflow/auto_mtf/valid_layouts_test.py b/mesh_tensorflow/auto_mtf/valid_layouts_test.py
index 47089455..4dd92214 100644
--- a/mesh_tensorflow/auto_mtf/valid_layouts_test.py
+++ b/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.
diff --git a/mesh_tensorflow/beam_search.py b/mesh_tensorflow/beam_search.py
index 90997228..c6a911f6 100644
--- a/mesh_tensorflow/beam_search.py
+++ b/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.
diff --git a/mesh_tensorflow/bert/__init__.py b/mesh_tensorflow/bert/__init__.py
index 898141c9..808b9020 100644
--- a/mesh_tensorflow/bert/__init__.py
+++ b/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.
diff --git a/mesh_tensorflow/bert/bert.py b/mesh_tensorflow/bert/bert.py
index aa4d6b3b..c42e6cdc 100644
--- a/mesh_tensorflow/bert/bert.py
+++ b/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.
diff --git a/mesh_tensorflow/bert/optimization.py b/mesh_tensorflow/bert/optimization.py
index 60d06360..361cf6ef 100644
--- a/mesh_tensorflow/bert/optimization.py
+++ b/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.
diff --git a/mesh_tensorflow/bert/run_classifier.py b/mesh_tensorflow/bert/run_classifier.py
index e4f072ab..9efb42de 100644
--- a/mesh_tensorflow/bert/run_classifier.py
+++ b/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.
diff --git a/mesh_tensorflow/bert/run_pretraining.py b/mesh_tensorflow/bert/run_pretraining.py
index 1778675f..b27b521a 100644
--- a/mesh_tensorflow/bert/run_pretraining.py
+++ b/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.
diff --git a/mesh_tensorflow/bert/run_squad.py b/mesh_tensorflow/bert/run_squad.py
index c1981eee..61ec2c0a 100644
--- a/mesh_tensorflow/bert/run_squad.py
+++ b/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.
diff --git a/mesh_tensorflow/bert/tokenization.py b/mesh_tensorflow/bert/tokenization.py
index 1133ed34..9ade5c13 100644
--- a/mesh_tensorflow/bert/tokenization.py
+++ b/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.
diff --git a/mesh_tensorflow/experimental/__init__.py b/mesh_tensorflow/experimental/__init__.py
index 898141c9..808b9020 100644
--- a/mesh_tensorflow/experimental/__init__.py
+++ b/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.
diff --git a/mesh_tensorflow/experimental/data_aug_lib.py b/mesh_tensorflow/experimental/data_aug_lib.py
index bb6e914e..b4e02956 100644
--- a/mesh_tensorflow/experimental/data_aug_lib.py
+++ b/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.
diff --git a/mesh_tensorflow/experimental/data_aug_lib_test.py b/mesh_tensorflow/experimental/data_aug_lib_test.py
index 0765726f..f3b94fb6 100644
--- a/mesh_tensorflow/experimental/data_aug_lib_test.py
+++ b/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.
diff --git a/mesh_tensorflow/experimental/input_reader.py b/mesh_tensorflow/experimental/input_reader.py
index 22565e76..fd0505bd 100644
--- a/mesh_tensorflow/experimental/input_reader.py
+++ b/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.
diff --git a/mesh_tensorflow/experimental/input_reader_test.py b/mesh_tensorflow/experimental/input_reader_test.py
index f6de01ee..60a3ebe7 100644
--- a/mesh_tensorflow/experimental/input_reader_test.py
+++ b/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.
diff --git a/mesh_tensorflow/experimental/model_executor.py b/mesh_tensorflow/experimental/model_executor.py
index 2ea5663f..78ee6364 100644
--- a/mesh_tensorflow/experimental/model_executor.py
+++ b/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.
diff --git a/mesh_tensorflow/experimental/offline_data_aug.py b/mesh_tensorflow/experimental/offline_data_aug.py
index 98721abd..aaa54329 100644
--- a/mesh_tensorflow/experimental/offline_data_aug.py
+++ b/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.
diff --git a/mesh_tensorflow/experimental/unet.py b/mesh_tensorflow/experimental/unet.py
index 17e6874b..34dfeb50 100644
--- a/mesh_tensorflow/experimental/unet.py
+++ b/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.
diff --git a/mesh_tensorflow/import_test.py b/mesh_tensorflow/import_test.py
index 3ea8563d..eddaea20 100644
--- a/mesh_tensorflow/import_test.py
+++ b/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.
diff --git a/mesh_tensorflow/layers.py b/mesh_tensorflow/layers.py
index 48bf4eb3..4270b723 100644
--- a/mesh_tensorflow/layers.py
+++ b/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.
diff --git a/mesh_tensorflow/layers_test.py b/mesh_tensorflow/layers_test.py
index 3eb935e6..c53b9b04 100644
--- a/mesh_tensorflow/layers_test.py
+++ b/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.
diff --git a/mesh_tensorflow/ops.py b/mesh_tensorflow/ops.py
index a7a6e786..7a5dd93a 100644
--- a/mesh_tensorflow/ops.py
+++ b/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.
diff --git a/mesh_tensorflow/ops_test.py b/mesh_tensorflow/ops_test.py
index 78c8ba6a..3551c7f3 100644
--- a/mesh_tensorflow/ops_test.py
+++ b/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.
diff --git a/mesh_tensorflow/ops_with_redefined_builtins.py b/mesh_tensorflow/ops_with_redefined_builtins.py
index 3b7c6a21..4fcefa99 100644
--- a/mesh_tensorflow/ops_with_redefined_builtins.py
+++ b/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.
diff --git a/mesh_tensorflow/optimize.py b/mesh_tensorflow/optimize.py
index a9c6023e..11147c02 100644
--- a/mesh_tensorflow/optimize.py
+++ b/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.
diff --git a/mesh_tensorflow/placement_mesh_impl.py b/mesh_tensorflow/placement_mesh_impl.py
index 76469ceb..68fd92be 100644
--- a/mesh_tensorflow/placement_mesh_impl.py
+++ b/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.
diff --git a/mesh_tensorflow/simd_mesh_impl.py b/mesh_tensorflow/simd_mesh_impl.py
index d5a4b72f..55874955 100644
--- a/mesh_tensorflow/simd_mesh_impl.py
+++ b/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.
diff --git a/mesh_tensorflow/simd_mesh_impl_test.py b/mesh_tensorflow/simd_mesh_impl_test.py
index 533ae546..95c89f88 100644
--- a/mesh_tensorflow/simd_mesh_impl_test.py
+++ b/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.
diff --git a/mesh_tensorflow/test_utils.py b/mesh_tensorflow/test_utils.py
index eda179e0..a17e3dd7 100644
--- a/mesh_tensorflow/test_utils.py
+++ b/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.
diff --git a/mesh_tensorflow/test_utils_test.py b/mesh_tensorflow/test_utils_test.py
index 8fa8469d..31d320ab 100644
--- a/mesh_tensorflow/test_utils_test.py
+++ b/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.
diff --git a/mesh_tensorflow/tpu_variables.py b/mesh_tensorflow/tpu_variables.py
index e5b1ba2d..6b0b23a8 100644
--- a/mesh_tensorflow/tpu_variables.py
+++ b/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.
diff --git a/mesh_tensorflow/transformer/__init__.py b/mesh_tensorflow/transformer/__init__.py
index 838ad4d6..a93a9319 100644
--- a/mesh_tensorflow/transformer/__init__.py
+++ b/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.
diff --git a/mesh_tensorflow/transformer/adaptive_softmax.py b/mesh_tensorflow/transformer/adaptive_softmax.py
index c2ce4676..43f15f09 100644
--- a/mesh_tensorflow/transformer/adaptive_softmax.py
+++ b/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.
diff --git a/mesh_tensorflow/transformer/adaptive_softmax_test.py b/mesh_tensorflow/transformer/adaptive_softmax_test.py
index 79db7fdc..078c8db3 100644
--- a/mesh_tensorflow/transformer/adaptive_softmax_test.py
+++ b/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.
diff --git a/mesh_tensorflow/transformer/attention.py b/mesh_tensorflow/transformer/attention.py
index 80a3fe0b..461e428e 100644
--- a/mesh_tensorflow/transformer/attention.py
+++ b/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.
diff --git a/mesh_tensorflow/transformer/dataset.py b/mesh_tensorflow/transformer/dataset.py
index e9e4aaa3..43f97e24 100644
--- a/mesh_tensorflow/transformer/dataset.py
+++ b/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.
diff --git a/mesh_tensorflow/transformer/dataset_test.py b/mesh_tensorflow/transformer/dataset_test.py
index 5bdf30c5..d34563a5 100644
--- a/mesh_tensorflow/transformer/dataset_test.py
+++ b/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.
diff --git a/mesh_tensorflow/transformer/evolved_transformer.py b/mesh_tensorflow/transformer/evolved_transformer.py
index 71532495..1a033674 100644
--- a/mesh_tensorflow/transformer/evolved_transformer.py
+++ b/mesh_tensorflow/transformer/evolved_transformer.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.
diff --git a/mesh_tensorflow/transformer/fixup_layers.py b/mesh_tensorflow/transformer/fixup_layers.py
index d6fc2511..cd45ebb8 100644
--- a/mesh_tensorflow/transformer/fixup_layers.py
+++ b/mesh_tensorflow/transformer/fixup_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.
diff --git a/mesh_tensorflow/transformer/funnel_transformer.py b/mesh_tensorflow/transformer/funnel_transformer.py
index d75c948e..8e35591a 100644
--- a/mesh_tensorflow/transformer/funnel_transformer.py
+++ b/mesh_tensorflow/transformer/funnel_transformer.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.
diff --git a/mesh_tensorflow/transformer/funnel_transformer_test.py b/mesh_tensorflow/transformer/funnel_transformer_test.py
index fff2fb9b..8da83f22 100644
--- a/mesh_tensorflow/transformer/funnel_transformer_test.py
+++ b/mesh_tensorflow/transformer/funnel_transformer_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.
diff --git a/mesh_tensorflow/transformer/gin/__init__.py b/mesh_tensorflow/transformer/gin/__init__.py
index 898141c9..808b9020 100644
--- a/mesh_tensorflow/transformer/gin/__init__.py
+++ b/mesh_tensorflow/transformer/gin/__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.
diff --git a/mesh_tensorflow/transformer/heterogeneous_moe.py b/mesh_tensorflow/transformer/heterogeneous_moe.py
index 53cae310..f2ae20b8 100644
--- a/mesh_tensorflow/transformer/heterogeneous_moe.py
+++ b/mesh_tensorflow/transformer/heterogeneous_moe.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.
diff --git a/mesh_tensorflow/transformer/learning_rate_schedules.py b/mesh_tensorflow/transformer/learning_rate_schedules.py
index 742bede9..363776f6 100644
--- a/mesh_tensorflow/transformer/learning_rate_schedules.py
+++ b/mesh_tensorflow/transformer/learning_rate_schedules.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.
diff --git a/mesh_tensorflow/transformer/learning_rate_schedules_test.py b/mesh_tensorflow/transformer/learning_rate_schedules_test.py
index 30e0159a..db43eb68 100644
--- a/mesh_tensorflow/transformer/learning_rate_schedules_test.py
+++ b/mesh_tensorflow/transformer/learning_rate_schedules_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.
diff --git a/mesh_tensorflow/transformer/main.py b/mesh_tensorflow/transformer/main.py
index e0e259d1..68de7214 100644
--- a/mesh_tensorflow/transformer/main.py
+++ b/mesh_tensorflow/transformer/main.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.
diff --git a/mesh_tensorflow/transformer/memory_layers.py b/mesh_tensorflow/transformer/memory_layers.py
index b62f0566..2325c900 100644
--- a/mesh_tensorflow/transformer/memory_layers.py
+++ b/mesh_tensorflow/transformer/memory_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.
diff --git a/mesh_tensorflow/transformer/memory_layers_test.py b/mesh_tensorflow/transformer/memory_layers_test.py
index 93ff185e..73b398fc 100644
--- a/mesh_tensorflow/transformer/memory_layers_test.py
+++ b/mesh_tensorflow/transformer/memory_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.
diff --git a/mesh_tensorflow/transformer/moe.py b/mesh_tensorflow/transformer/moe.py
index 910adf04..442729eb 100644
--- a/mesh_tensorflow/transformer/moe.py
+++ b/mesh_tensorflow/transformer/moe.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.
diff --git a/mesh_tensorflow/transformer/t2t_vocabulary.py b/mesh_tensorflow/transformer/t2t_vocabulary.py
index 37f3c08e..bdae8382 100644
--- a/mesh_tensorflow/transformer/t2t_vocabulary.py
+++ b/mesh_tensorflow/transformer/t2t_vocabulary.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.
diff --git a/mesh_tensorflow/transformer/transformer.py b/mesh_tensorflow/transformer/transformer.py
index 7e771f85..e6a3d11c 100644
--- a/mesh_tensorflow/transformer/transformer.py
+++ b/mesh_tensorflow/transformer/transformer.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.
diff --git a/mesh_tensorflow/transformer/transformer_layers.py b/mesh_tensorflow/transformer/transformer_layers.py
index 34ab2033..4be16769 100644
--- a/mesh_tensorflow/transformer/transformer_layers.py
+++ b/mesh_tensorflow/transformer/transformer_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.
diff --git a/mesh_tensorflow/transformer/transformer_layers_test.py b/mesh_tensorflow/transformer/transformer_layers_test.py
index 805f24d2..9b5e35d9 100644
--- a/mesh_tensorflow/transformer/transformer_layers_test.py
+++ b/mesh_tensorflow/transformer/transformer_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.
diff --git a/mesh_tensorflow/transformer/universal_transformer.py b/mesh_tensorflow/transformer/universal_transformer.py
index 55a1bd8d..ee57a445 100644
--- a/mesh_tensorflow/transformer/universal_transformer.py
+++ b/mesh_tensorflow/transformer/universal_transformer.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.
diff --git a/mesh_tensorflow/transformer/utils.py b/mesh_tensorflow/transformer/utils.py
index 017acf55..552ba0f7 100644
--- a/mesh_tensorflow/transformer/utils.py
+++ b/mesh_tensorflow/transformer/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.
@@ -48,8 +48,8 @@
 
 from tensorflow.core.protobuf import rewriter_config_pb2  # pylint: disable=g-direct-tensorflow-import
 from tensorflow.python.ops import resources  # pylint: disable=g-direct-tensorflow-import
-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.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
 
 try:
   tf.flags.DEFINE_multi_string("gin_file", None, "Path to a Gin file.")
diff --git a/mesh_tensorflow/transformer/utils_test.py b/mesh_tensorflow/transformer/utils_test.py
index c6d2fc1c..0e7d5116 100644
--- a/mesh_tensorflow/transformer/utils_test.py
+++ b/mesh_tensorflow/transformer/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.
diff --git a/mesh_tensorflow/transformer/vocab_embeddings.py b/mesh_tensorflow/transformer/vocab_embeddings.py
index 89c096fb..b27ebc06 100644
--- a/mesh_tensorflow/transformer/vocab_embeddings.py
+++ b/mesh_tensorflow/transformer/vocab_embeddings.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.
diff --git a/mesh_tensorflow/transformer/vocab_embeddings_test.py b/mesh_tensorflow/transformer/vocab_embeddings_test.py
index 0c063384..da8f2f90 100644
--- a/mesh_tensorflow/transformer/vocab_embeddings_test.py
+++ b/mesh_tensorflow/transformer/vocab_embeddings_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.
diff --git a/mesh_tensorflow/transformer/vocabulary.py b/mesh_tensorflow/transformer/vocabulary.py
index 1f156d86..13f7d780 100644
--- a/mesh_tensorflow/transformer/vocabulary.py
+++ b/mesh_tensorflow/transformer/vocabulary.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.
diff --git a/mesh_tensorflow/utils.py b/mesh_tensorflow/utils.py
index 040fd7d5..f25683a7 100644
--- a/mesh_tensorflow/utils.py
+++ b/mesh_tensorflow/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.
diff --git a/mesh_tensorflow/utils_test.py b/mesh_tensorflow/utils_test.py
index 71025217..65b485dc 100644
--- a/mesh_tensorflow/utils_test.py
+++ b/mesh_tensorflow/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.