Update tensorflow to version 2.16.1.

RELEASE.md

@@ -8,6 +8,10 @@ and this project adheres to
 
 ## Unreleased
 
+### Changed
+
+*   Updated `tensorflow` to version `2.16.1`.
+
 ## Release 0.81.0
 
 ### Added

WORKSPACE

@@ -23,12 +23,15 @@ http_archive(
     strip_prefix = "bazel-skylib-1.3.0",
 )
 
-http_archive(
-    name = "com_github_grpc_grpc",
-    url = "https://github.com/grpc/grpc/archive/refs/tags/v1.50.0.tar.gz",
-    sha256 = "76900ab068da86378395a8e125b5cc43dfae671e09ff6462ddfef18676e2165a",
-    strip_prefix = "grpc-1.50.0",
-)
+# TODO: b/347937669 - Temporarily disable the direct dependency on
+# `com_github_grpc_grpc`, for now we pick this dependency up from the
+# TensorFlow workspace.
+# http_archive(
+#     name = "com_github_grpc_grpc",
+#     url = "https://github.com/grpc/grpc/archive/refs/tags/v1.50.0.tar.gz",
+#     sha256 = "76900ab068da86378395a8e125b5cc43dfae671e09ff6462ddfef18676e2165a",
+#     strip_prefix = "grpc-1.50.0",
+# )
 
 http_archive(
     name = "com_google_benchmark",
@@ -74,8 +77,8 @@ http_archive(
     strip_prefix = "protobuf-3.21.9",
 )
 
-# TODO: b/333391041 - Temporarily disable the direct dependency on
-# `eigen`, for now we pick this dependency up from the TensorFlow workspace.
+# TODO: b/333391041 - Temporarily disable the direct dependency on `eigen`, for
+# now we pick this dependency up from the TensorFlow workspace.
 # http_archive(
 #     name = "eigen",
 #     url = "https://gitlab.com/libeigen/eigen/archive/refs/tags/3.4.0.tar.gz",
@@ -91,14 +94,14 @@ http_archive(
 # https://github.com/google-parfait/tensorflow-federated/blob/main/requirements.txt.
 http_archive(
     name = "org_tensorflow",
-    url = "https://github.com/tensorflow/tensorflow/archive/refs/tags/v2.14.0.tar.gz",
+    url = "https://github.com/tensorflow/tensorflow/archive/refs/tags/v2.16.1.tar.gz",
     patches = [
         "//third_party/tensorflow:internal_visibility.patch",
         "//third_party/tensorflow:python_toolchain.patch",
         "//third_party/tensorflow:tf2xla_visibility.patch",
     ],
-    sha256 = "ce357fd0728f0d1b0831d1653f475591662ec5bca736a94ff789e6b1944df19f",
-    strip_prefix = "tensorflow-2.14.0",
+    sha256 = "c729e56efc945c6df08efe5c9f5b8b89329c7c91b8f40ad2bb3e13900bd4876d",
+    strip_prefix = "tensorflow-2.16.1",
 )
 
 # This commit is determined by
@@ -140,6 +143,24 @@ http_archive(
     strip_prefix = "rules_python-0.31.0",
 )
 
+# TODO: b/347940678 - Temporarily disable the direct dependency on `tsl`, for
+# now we pick this dependency up from the TensorFlow workspace.
+# http_archive(
+#     name = "tsl",
+#     url = "https://github.com/google/tsl/archive/5fdb99d897eadddd136c41156fbda82c6244d60a.tar.gz",
+#     sha256 = "414f9929d6a631aa225093e24e657e1a28c5e84da8e280f1636db65fe24e4b83",
+#     strip_prefix = "tsl-5fdb99d897eadddd136c41156fbda82c6244d60a",
+# )
+
+# TODO: b/347937714 - Temporarily disable the direct dependency on `xla`, for
+# now we pick this dependency up from the TensorFlow workspace.
+# http_archive(
+#     name = "xla",
+#     url = "https://github.com/openxla/xla/archive/940e3a27542b7ce76666173e7b287aa2a9263916.tar.gz",
+#     sha256 = "bcdc778e5a456839869dea796117b723bdea488075bd9555fe118fd8d6fcf25e",
+#     strip_prefix = "xla-940e3a27542b7ce76666173e7b287aa2a9263916",
+# )
+
 #
 # Inlined transitive dependencies, grouped by direct dependency.
 #
@@ -153,14 +174,17 @@ http_archive(
     build_file = "@pybind11_bazel//:pybind11.BUILD",
 )
 
+# TODO: b/347937669 - Temporarily disable the direct dependency on
+# `com_github_grpc_grpc`, for now we pick this dependency up from the
+# TensorFlow workspace.
 # Required by com_github_grpc_grpc. This commit is determined by
 # https://github.com/grpc/grpc/blob/v1.50.0/bazel/grpc_deps.bzl#L344.
-http_archive(
-    name = "upb",
-    url = "https://github.com/protocolbuffers/upb/archive/e4635f223e7d36dfbea3b722a4ca4807a7e882e2.tar.gz",
-    sha256 = "017a7e8e4e842d01dba5dc8aa316323eee080cd1b75986a7d1f94d87220e6502",
-    strip_prefix = "upb-e4635f223e7d36dfbea3b722a4ca4807a7e882e2",
-)
+# http_archive(
+#     name = "upb",
+#     url = "https://github.com/protocolbuffers/upb/archive/e4635f223e7d36dfbea3b722a4ca4807a7e882e2.tar.gz",
+#     sha256 = "017a7e8e4e842d01dba5dc8aa316323eee080cd1b75986a7d1f94d87220e6502",
+#     strip_prefix = "upb-e4635f223e7d36dfbea3b722a4ca4807a7e882e2",
+# )
 
 #
 # Transitive dependencies, grouped by direct dependency.

pyproject.toml

@@ -46,7 +46,7 @@ dependencies = [
   'scipy~=1.9.3',
   'tensorflow-model-optimization==0.7.5',
   'tensorflow-privacy==0.9.0',
-  'tensorflow>=2.14.0,==2.14.*',
+  'tensorflow~=2.16.1',
   'tqdm~=4.64',
   'typing-extensions>=4.5.0,==4.5.*',
   # TODO: b/315515548 - Required because current dependencies are pulling in

requirements.txt

@@ -44,7 +44,7 @@ typing-extensions>=4.5.0,==4.5.*
 
 # The version of this dependency should match the version in
 # https://github.com/google-parfait/tensorflow-federated/blob/main/WORKSPACE.
-tensorflow>=2.14.0,==2.14.*
+tensorflow~=2.16.1
 
 # TODO: b/315515548 - Required because current dependencies are pulling in later
 # versions which are causing opaque failures.

tensorflow_federated/cc/core/impl/aggregation/tensorflow/BUILD

@@ -29,10 +29,10 @@ cc_library(
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
+        "@local_tsl//tsl/platform:refcount",
         "@org_tensorflow//tensorflow/core:framework",
         "@org_tensorflow//tensorflow/core:protos_all_cc",
         "@org_tensorflow//tensorflow/core/platform:tstring",
-        "@org_tensorflow//tensorflow/tsl/platform:refcount",
     ],
 )
 
@@ -130,8 +130,8 @@ cc_library(
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:cord",
+        "@local_tsl//tsl/platform:env",
         "@org_tensorflow//tensorflow/core/platform:env",
-        "@org_tensorflow//tensorflow/tsl/platform:env",
     ],
 )
 
@@ -165,8 +165,8 @@ cc_library(
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:cord",
+        "@local_tsl//tsl/platform:env",
         "@org_tensorflow//tensorflow/core/platform:env",
-        "@org_tensorflow//tensorflow/tsl/platform:env",
     ],
 )
 

tensorflow_federated/cc/core/impl/aggregation/tensorflow/converters.cc

@@ -29,7 +29,7 @@
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/tensor_shape.pb.h"
 #include "tensorflow/core/platform/tstring.h"
-#include "tensorflow/tsl/platform/refcount.h"
+#include "tsl/platform/refcount.h"
 #include "tensorflow_federated/cc/core/impl/aggregation/base/monitoring.h"
 #include "tensorflow_federated/cc/core/impl/aggregation/core/datatype.h"
 #include "tensorflow_federated/cc/core/impl/aggregation/core/tensor.h"

tensorflow_federated/cc/core/impl/aggregation/tensorflow/tensorflow_checkpoint_builder_factory.cc

@@ -29,7 +29,7 @@
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "tensorflow/core/platform/file_system.h"
-#include "tensorflow/tsl/platform/env.h"
+#include "tsl/platform/env.h"
 #include "tensorflow_federated/cc/core/impl/aggregation/base/monitoring.h"
 #include "tensorflow_federated/cc/core/impl/aggregation/core/tensor.h"
 #include "tensorflow_federated/cc/core/impl/aggregation/protocol/checkpoint_builder.h"

tensorflow_federated/cc/core/impl/aggregation/tensorflow/tensorflow_checkpoint_parser_factory.cc

@@ -29,7 +29,7 @@
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "tensorflow/core/platform/file_system.h"
-#include "tensorflow/tsl/platform/env.h"
+#include "tsl/platform/env.h"
 #include "tensorflow_federated/cc/core/impl/aggregation/base/monitoring.h"
 #include "tensorflow_federated/cc/core/impl/aggregation/core/tensor.h"
 #include "tensorflow_federated/cc/core/impl/aggregation/protocol/checkpoint_parser.h"

tensorflow_federated/cc/core/impl/executors/BUILD

@@ -342,7 +342,6 @@ pybind_extension(
         "@org_tensorflow//tensorflow/core:framework",
         "@org_tensorflow//tensorflow/core/platform:status",
         "@org_tensorflow//tensorflow/core/platform:strcat",
-        "@org_tensorflow//tensorflow/python/lib/core:ndarray_tensor_headers",
         "@pybind11_abseil//pybind11_abseil:absl_casters",
         "@pybind11_abseil//pybind11_abseil:status_casters",
         "@pybind11_protobuf//pybind11_protobuf:native_proto_caster",
@@ -551,7 +550,6 @@ cc_library(
         "//tensorflow_federated/proto/v0:executor_cc_grpc_proto",
         "//tensorflow_federated/proto/v0:executor_cc_proto",
         "@com_github_grpc_grpc//:grpc++",
-        "@com_github_grpc_grpc//:grpc_security_base",
         "@com_google_absl//absl/strings",
         "@com_google_googletest//:gtest",
     ],
@@ -1114,20 +1112,20 @@ cc_library(
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",
+        "@local_xla//xla:literal",
+        "@local_xla//xla:shape_util",
+        "@local_xla//xla:xla_data_proto_cc",
+        "@local_xla//xla:xla_proto_cc",
+        "@local_xla//xla/client",
+        "@local_xla//xla/client:client_library",
+        "@local_xla//xla/client:global_data",
+        "@local_xla//xla/client:xla_computation",
+        "@local_xla//xla/service:hlo_proto_cc",
+        "@local_xla//xla/stream_executor",
+        "@local_xla//xla/stream_executor:multi_platform_manager",
+        "@local_xla//xla/stream_executor/host:host_platform",  # buildcleaner: keep # Linking in the host platform here ensures that the stream executor can execute on CPU.
         "@org_tensorflow//tensorflow/compiler/jit:xla_cpu_jit",  # buildcleaner: keep # Linking in this dependency ensures that XLA can compile its code for the CPU host.
         "@org_tensorflow//tensorflow/compiler/tf2xla:common",
-        "@org_tensorflow//tensorflow/compiler/xla:literal",
-        "@org_tensorflow//tensorflow/compiler/xla:shape_util",
-        "@org_tensorflow//tensorflow/compiler/xla:xla_data_proto_cc",
-        "@org_tensorflow//tensorflow/compiler/xla:xla_proto_cc",
-        "@org_tensorflow//tensorflow/compiler/xla/client",
-        "@org_tensorflow//tensorflow/compiler/xla/client:client_library",
-        "@org_tensorflow//tensorflow/compiler/xla/client:global_data",
-        "@org_tensorflow//tensorflow/compiler/xla/client:xla_computation",
-        "@org_tensorflow//tensorflow/compiler/xla/service:hlo_proto_cc",
-        "@org_tensorflow//tensorflow/compiler/xla/stream_executor",
-        "@org_tensorflow//tensorflow/compiler/xla/stream_executor:multi_platform_manager",
-        "@org_tensorflow//tensorflow/compiler/xla/stream_executor/host:host_platform",  # buildcleaner: keep # Linking in the host platform here ensures that the stream executor can execute on CPU.
         "@org_tensorflow//tensorflow/core:framework",
         "@org_tensorflow//tensorflow/core:protos_all_cc",
     ],
@@ -1156,13 +1154,13 @@ tff_cc_cpu_gpu_test(
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
+        "@local_xla//xla:shape_util",
+        "@local_xla//xla:xla_data_proto_cc",
+        "@local_xla//xla/client:xla_builder",
+        "@local_xla//xla/client:xla_computation",
+        "@local_xla//xla/service:platform_util",
+        "@local_xla//xla/stream_executor",
         "@org_tensorflow//tensorflow/compiler/tf2xla:common",
-        "@org_tensorflow//tensorflow/compiler/xla:shape_util",
-        "@org_tensorflow//tensorflow/compiler/xla:xla_data_proto_cc",
-        "@org_tensorflow//tensorflow/compiler/xla/client:xla_builder",
-        "@org_tensorflow//tensorflow/compiler/xla/client:xla_computation",
-        "@org_tensorflow//tensorflow/compiler/xla/service:platform_util",
-        "@org_tensorflow//tensorflow/compiler/xla/stream_executor",
         "@org_tensorflow//tensorflow/core:framework",
         "@org_tensorflow//tensorflow/core:protos_all_cc",
         "@org_tensorflow//tensorflow/core/platform:statusor",
@@ -1193,15 +1191,15 @@ cc_test(
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
+        "@local_xla//xla:shape_util",
+        "@local_xla//xla:xla_data_proto_cc",
+        "@local_xla//xla/client:xla_builder",
+        "@local_xla//xla/client:xla_computation",
+        "@local_xla//xla/service:platform_util",
+        "@local_xla//xla/stream_executor",
+        "@local_xla//xla/stream_executor/cuda:cuda_platform",  # buildcleaner: keep # Linking in the host platform here ensures that the stream executor can execute on GPU.
         "@org_tensorflow//tensorflow/compiler/jit:xla_gpu_jit",  # buildcleaner: keep # Linking in this dependency ensures that XLA can compile its code for the GPU.
         "@org_tensorflow//tensorflow/compiler/tf2xla:common",
-        "@org_tensorflow//tensorflow/compiler/xla:shape_util",
-        "@org_tensorflow//tensorflow/compiler/xla:xla_data_proto_cc",
-        "@org_tensorflow//tensorflow/compiler/xla/client:xla_builder",
-        "@org_tensorflow//tensorflow/compiler/xla/client:xla_computation",
-        "@org_tensorflow//tensorflow/compiler/xla/service:platform_util",
-        "@org_tensorflow//tensorflow/compiler/xla/stream_executor",
-        "@org_tensorflow//tensorflow/compiler/xla/stream_executor/cuda:cuda_platform",  # buildcleaner: keep # Linking in the host platform here ensures that the stream executor can execute on GPU.
         "@org_tensorflow//tensorflow/core:framework",
         "@org_tensorflow//tensorflow/core:protos_all_cc",
         "@org_tensorflow//tensorflow/core/platform:statusor",
@@ -1220,9 +1218,9 @@ cc_library(
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
-        "@org_tensorflow//tensorflow/compiler/xla:literal",
-        "@org_tensorflow//tensorflow/compiler/xla:shape_util",
-        "@org_tensorflow//tensorflow/compiler/xla:types",
+        "@local_xla//xla:literal",
+        "@local_xla//xla:shape_util",
+        "@local_xla//xla:types",
     ],
 )
 
@@ -1241,10 +1239,10 @@ cc_test(
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@eigen_archive//:eigen3",
-        "@org_tensorflow//tensorflow/compiler/xla:literal",
-        "@org_tensorflow//tensorflow/compiler/xla:literal_util",
-        "@org_tensorflow//tensorflow/compiler/xla:shape_util",
-        "@org_tensorflow//tensorflow/compiler/xla:types",
-        "@org_tensorflow//tensorflow/compiler/xla:xla_data_proto_cc",
+        "@local_xla//xla:literal",
+        "@local_xla//xla:literal_util",
+        "@local_xla//xla:shape_util",
+        "@local_xla//xla:types",
+        "@local_xla//xla:xla_data_proto_cc",
     ],
 )

tensorflow_federated/cc/core/impl/executors/array_test_utils.h

@@ -23,7 +23,7 @@ limitations under the License
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
-#include "third_party/eigen3/Eigen/Core"
+#include "Eigen/Core"
 #include "tensorflow_federated/proto/v0/array.pb.h"
 #include "tensorflow_federated/proto/v0/data_type.pb.h"
 

tensorflow_federated/cc/core/impl/executors/tensorflow_utils.cc

@@ -24,7 +24,7 @@ limitations under the License
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
-#include "third_party/eigen3/Eigen/Core"
+#include "Eigen/Core"
 #include "tensorflow/core/framework/attr_value.pb.h"
 #include "tensorflow/core/framework/graph.pb.h"
 #include "tensorflow/core/framework/node_def.pb.h"

tensorflow_federated/cc/core/impl/executors/tensorflow_utils_test.cc

@@ -22,7 +22,7 @@ limitations under the License
 #include "googletest/include/gtest/gtest.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "third_party/eigen3/Eigen/Core"
+#include "Eigen/Core"
 #include "tensorflow/core/framework/numeric_types.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/tensor_testutil.h"

tensorflow_federated/cc/core/impl/executors/xla_executor.cc

@@ -33,16 +33,16 @@ limitations under the License
 #include "absl/types/span.h"
 #include "tensorflow/compiler/tf2xla/literal_util.h"
 #include "tensorflow/compiler/tf2xla/type_util.h"
-#include "tensorflow/compiler/xla/client/client.h"
-#include "tensorflow/compiler/xla/client/client_library.h"
-#include "tensorflow/compiler/xla/client/global_data.h"
-#include "tensorflow/compiler/xla/client/xla_computation.h"
-#include "tensorflow/compiler/xla/literal.h"
-#include "tensorflow/compiler/xla/service/hlo.pb.h"
-#include "tensorflow/compiler/xla/shape.h"
-#include "tensorflow/compiler/xla/stream_executor/platform.h"
-#include "tensorflow/compiler/xla/xla.pb.h"
-#include "tensorflow/compiler/xla/xla_data.pb.h"
+#include "xla/client/client.h"
+#include "xla/client/client_library.h"
+#include "xla/client/global_data.h"
+#include "xla/client/xla_computation.h"
+#include "xla/literal.h"
+#include "xla/service/hlo.pb.h"
+#include "xla/shape.h"
+#include "xla/stream_executor/platform.h"
+#include "xla/xla.pb.h"
+#include "xla/xla_data.pb.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor.pb.h"
 #include "tensorflow/core/framework/tensor_shape.h"
@@ -58,7 +58,7 @@ limitations under the License
 // clang-format off
 // In TF 2.17 MultiPlatformManager was renamed to PlatformManager. Remove
 // this code when the OSS build gets updated to TF 2.17+.
-#include "tensorflow/compiler/xla/stream_executor/multi_platform_manager.h"
+#include "xla/stream_executor/multi_platform_manager.h"
 namespace stream_executor {
 using PlatformManager = MultiPlatformManager;
 } // namespace stream_executor

tensorflow_federated/cc/core/impl/executors/xla_executor_test.cc

@@ -32,13 +32,13 @@ limitations under the License
 #include "absl/status/statusor.h"
 #include "tensorflow/compiler/tf2xla/shape_util.h"
 #include "tensorflow/compiler/tf2xla/type_util.h"
-#include "tensorflow/compiler/xla/client/xla_builder.h"
-#include "tensorflow/compiler/xla/client/xla_computation.h"
-#include "tensorflow/compiler/xla/service/platform_util.h"
-#include "tensorflow/compiler/xla/shape.h"
-#include "tensorflow/compiler/xla/shape_util.h"
-#include "tensorflow/compiler/xla/stream_executor/platform.h"
-#include "tensorflow/compiler/xla/xla_data.pb.h"
+#include "xla/client/xla_builder.h"
+#include "xla/client/xla_computation.h"
+#include "xla/service/platform_util.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
+#include "xla/stream_executor/platform.h"
+#include "xla/xla_data.pb.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/platform/statusor.h"