[pull] master from tensorflow:master

.bazelrc

@@ -709,6 +709,7 @@ common:rbe_win_base --nobuild_runfile_links
 common:rbe_win_base --define=override_eigen_strong_inline=true
 
 common:rbe_windows_x86_cpu_2022 --config=rbe_win_base --config=windows_x86_cpu_2022
+common:rbe_windows_x86_cpu_2022 --features=-with_wrapper
 
 # Override the remote instance and BES instance names to use the new RBE pools
 common:rbe_windows_x86_cpu_2022 --remote_instance_name=projects/ml-oss-rbe-general/instances/default_instance

.github/workflows/pylint-presubmit.yml

@@ -44,7 +44,7 @@ jobs:
     - name: Install Python dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install pylint==2.13.9 numpy wheel
+        pip install pylint==2.13.9 "numpy<2.0.0" wheel
     - name: Run PyLint on changed files
       run: |
         echo "${{ steps.get_file_changes.outputs.files}}" | tr " " "\n" | grep ".py$" | xargs pylint --rcfile=tensorflow/tools/ci_build/pylintrc

tensorflow/c/tf_datatype.h

@@ -65,7 +65,9 @@ typedef enum TF_DataType {
   TF_UINT4 = 30,
   TF_INT2 = 31,
   TF_UINT2 = 32,
-  TF_FLOAT4_E2M1FN = 33  // 2 exponent bits, 1 mantissa bit, finite-only
+  TF_FLOAT4_E2M1FN = 33,  // 2 exponent bits, 1 mantissa bit, finite-only
+  TF_FLOAT8_E8M0FNU = 34  // 8 exponent bits, 0 mantissa bits, finite-only, with
+                          // NaN.
 } TF_DataType;
 
 // TF_DataTypeSize returns the sizeof() for the underlying type corresponding

tensorflow/cc/saved_model/BUILD

@@ -722,6 +722,8 @@ tf_cc_test(
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_googletest//:gtest",
+        "@xla//xla/tsl/lib/core:status_test_util",
+        "@xla//xla/tsl/platform:status_matchers",
     ],
 )
 

tensorflow/cc/saved_model/fingerprinting.cc

@@ -174,6 +174,11 @@ absl::StatusOr<FingerprintDef> CreateFingerprintDefPb(
   SavedModel saved_model;
   TF_RETURN_IF_ERROR(ReadBinaryProto(Env::Default(), pb_file, &saved_model));
 
+  if (saved_model.meta_graphs_size() == 0) {
+    return absl::InvalidArgumentError(
+        "SavedModel (.pb) contains no MetaGraphs.");
+  }
+
   // Create a copy of `metagraph` which will be used and mutated for fingerprint
   // computation.
   FingerprintDef fingerprint_def;

tensorflow/cc/saved_model/fingerprinting_test.cc

@@ -17,12 +17,15 @@ limitations under the License.
 
 #include <string>
 
+#include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/numeric/int128.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/numbers.h"
 #include "absl/strings/string_view.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/tsl/platform/status_matchers.h"
 #include "tensorflow/core/framework/graph.pb.h"
 #include "tensorflow/core/framework/versions.pb.h"
 #include "tensorflow/core/platform/env.h"
@@ -206,5 +209,19 @@ TEST(FingerprintingTest, TestSingleprint) {
             const_singleprint);
 }
 
+TEST(FingerprintingTest, CreateFingerprintDefPbEmptyMetaGraphsReturnsError) {
+  const std::string model_dir =
+      io::JoinPath(::testing::TempDir(), "empty_metagraphs_model");
+  TF_ASSERT_OK(Env::Default()->RecursivelyCreateDir(model_dir));
+  const std::string pb_file = io::JoinPath(model_dir, "saved_model.pb");
+  SavedModel saved_model;
+  TF_ASSERT_OK(WriteBinaryProto(Env::Default(), pb_file, saved_model));
+
+  EXPECT_THAT(
+      CreateFingerprintDef(model_dir).status(),
+      tsl::testing::StatusIs(absl::StatusCode::kInvalidArgument,
+                             "SavedModel (.pb) contains no MetaGraphs."));
+}
+
 }  // namespace
 }  // namespace tensorflow::saved_model::fingerprinting

tensorflow/compiler/aot/BUILD

@@ -148,6 +148,7 @@ cc_library(
         "@xla//xla/client:client_library",
         "@xla//xla/client:compile_only_client",
         "@xla//xla/hlo/builder:xla_computation",
+        "@xla//xla/pjrt:pjrt_compiler",
         "@xla//xla/service:compiler",
         "@xla//xla/service/cpu:cpu_aot_compilation_result",
         "@xla//xla/service/cpu:cpu_compiler",

tensorflow/compiler/aot/compile.cc

@@ -42,6 +42,7 @@ limitations under the License.
 #include "xla/client/client_library.h"
 #include "xla/client/compile_only_client.h"
 #include "xla/hlo/builder/xla_computation.h"
+#include "xla/pjrt/pjrt_compiler.h"
 #include "xla/service/compiler.h"
 #include "xla/service/cpu/cpu_aot_compilation_result.h"
 #include "xla/shape.h"
@@ -174,8 +175,10 @@ absl::Status CompileGraph(GraphDef graph_def, const tf2xla::Config& config,
         graph_def, config, &computation, flags.debug_info,
         flags.debug_info_path_begin_marker));
   } else {
+    TF_ASSIGN_OR_RETURN(auto* pjrt_compiler,
+                        xla::GetDefaultPjRtCompiler(xla::CpuName()));
     TF_RETURN_IF_ERROR(ConvertGraphDefToXla(std::move(graph_def), config,
-                                            client, &computation));
+                                            pjrt_compiler, &computation));
   }
 
   if (flags.experimental_quantize && *quantize_xla) {

tensorflow/compiler/jit/BUILD

@@ -156,6 +156,7 @@ cc_library(
         "//tensorflow/compiler/tf2xla/kernels:xla_ops",
         "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:lib",
+        "//tensorflow/core/tfrt/common:pjrt_cpu_client_registration",
         "@com_google_absl//absl/memory",
         "@xla//xla/stream_executor:platform_manager",
     ] + if_libtpu(
@@ -291,7 +292,6 @@ cc_library(
     # Public visibility is needed for external TF/XLA backends.
     visibility = ["//visibility:public"],
     deps = [
-        ":common",
         ":device_compilation_cluster_signature",
         ":device_compilation_profiler",
         ":device_compiler",
@@ -304,13 +304,11 @@ cc_library(
         ":variable_info_util",
         ":xla_compile_util",
         ":xla_device_compiler_client",
-        ":xla_device_context",
         ":xla_launch_util",
         ":xla_tensor",
         "//tensorflow/compiler/jit/ops:xla_ops",
         "//tensorflow/compiler/tf2xla:common",
         "//tensorflow/compiler/tf2xla:layout_util",
-        "//tensorflow/compiler/tf2xla:tf2xla_util",
         "//tensorflow/compiler/tf2xla:xla_compiler",
         "//tensorflow/compiler/tf2xla:xla_op_registry",
         "//tensorflow/compiler/tf2xla/kernels:xla_dummy_ops",
@@ -346,36 +344,28 @@ cc_library(
         "//tensorflow/core/kernels/data:options_dataset_op",
         "//tensorflow/core/kernels/data:prefetch_dataset_op",
         "//tensorflow/core/platform:stream_executor_no_cuda",
-        "//tensorflow/core/profiler/lib:traceme",
         "//tensorflow/core/tfrt/common:async_value_tensor",
         "//tensorflow/core/tfrt/common:create_pjrt_client_util",
         "//tensorflow/core/tfrt/common:global_state",
         "//tensorflow/core/tfrt/common:pjrt_util",
         "//tensorflow/core/tpu:tpu_defs",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base",
-        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
-        "@com_google_absl//absl/memory",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:optional",
         "@com_google_absl//absl/types:span",
-        "@xla//xla:executable_run_options",
+        "@xla//xla:shape_util",
         "@xla//xla/client:client_library",
         "@xla//xla/client:local_client",
-        "@xla//xla/hlo/ir:hlo",
         "@xla//xla/pjrt:pjrt_client",
-        "@xla//xla/pjrt:tf_pjrt_client",
         "@xla//xla/service:compiler",
-        "@xla//xla/service:executable",
-        "@xla//xla/service:stream_pool",
-        "@xla//xla/service/gpu:gpu_executable_run_options",
+        "@xla//xla/service:shaped_buffer",
+        "@xla//xla/stream_executor:platform",
         "@xla//xla/stream_executor:platform_manager",
-        "@xla//xla/stream_executor:stream",
-        "@xla//xla/stream_executor/host:host_platform_id",
         "@xla//xla/stream_executor/integrations:tf_allocator_adapter",
         "@xla//xla/tsl/framework:device_id_utils",
     ],
@@ -403,11 +393,9 @@ cc_library(
     deps = [
         ":device_compilation_cluster_signature",
         ":device_compilation_profiler",
+        ":device_compiler",
         ":jit_compilation_passes",
-        ":variable_info",
-        ":variable_info_util",
         ":xla_device_no_jit_rewrite_registration",
-        ":xla_launch_util",
         ":xla_tensor",
         "//tensorflow/compiler/jit/ops:xla_ops",
         "//tensorflow/compiler/tf2xla:layout_util",
@@ -446,6 +434,7 @@ cc_library(
         "//tensorflow/core/kernels/data:prefetch_dataset_op",
         "//tensorflow/core/platform:stream_executor_no_cuda",
         "//tensorflow/core/tfrt/common:async_value_tensor",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/types:optional",
         "@xla//xla/client:local_client",
         "@xla//xla/pjrt:pjrt_client",
@@ -476,13 +465,12 @@ cc_library(
         "//tensorflow/compiler/mlir/tensorflow:dump_graph",
         "//tensorflow/core:framework_internal",
         "//tensorflow/core:framework_types_hdr",
-        "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/tpu/kernels:sparse_core_xla_flags_defaults",
         "@com_google_absl//absl/base",
         "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/log",
         "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/types:optional",
         "@xla//xla:parse_flags_from_env",
     ],
 )
@@ -496,10 +484,8 @@ cc_library(
     deps = [
         "//tensorflow/core:framework_internal",
         "//tensorflow/core:framework_types_hdr",
-        "//tensorflow/core:lib",
         "//tensorflow/core/protobuf:for_core_protos_cc",
         "@com_google_absl//absl/container:flat_hash_set",
-        "@com_google_absl//absl/types:optional",
     ],
 )
 
@@ -629,7 +615,6 @@ cc_library(
         ":pjrt_tensor_buffer_util",
         ":variable_info",
         ":variable_info_util",
-        ":xla_tensor",
         "//tensorflow/compiler/tf2xla:common",
         "//tensorflow/compiler/tf2xla:xla_compiler",
         "//tensorflow/compiler/tf2xla:xla_resource",
@@ -644,6 +629,7 @@ cc_library(
         "//tensorflow/core/common_runtime/gpu:gpu_serving_device_selector",
         "//tensorflow/core/tfrt/common:async_value_tensor",
         "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/base",
         "@com_google_absl//absl/cleanup",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
@@ -660,18 +646,13 @@ cc_library(
         "@xla//xla/client:local_client",
         "@xla//xla/hlo/ir:hlo",
         "@xla//xla/pjrt:pjrt_client",
-        "@xla//xla/pjrt:pjrt_common",
+        "@xla//xla/pjrt:pjrt_compiler",
         "@xla//xla/pjrt:pjrt_executable",
+        "@xla//xla/pjrt/plugin/xla_cpu:cpu_execute_options",
+        "@xla//xla/runtime:device_id",
         "@xla//xla/service:executable",
-        "@xla//xla/service:maybe_owning_device_memory",
         "@xla//xla/service:shaped_buffer",
-        "@xla//xla/service:transfer_manager",
-        "@xla//xla/stream_executor:device_memory",
-        "@xla//xla/stream_executor:device_memory_allocator",
-        "@xla//xla/stream_executor:event",
-        "@xla//xla/stream_executor:platform",
-        "@xla//xla/stream_executor:platform_manager",
-        "@xla//xla/stream_executor/host:host_platform_id",
+        "@xla//xla/stream_executor:device_address_allocator",
         "@xla//xla/tsl/framework:device_id_utils",
         "@xla//xla/tsl/framework:serving_device_selector_policies",
     ],
@@ -875,30 +856,41 @@ cc_library(
     visibility = [":internal"],
     deps = [
         ":compilability_check_util",
+        ":device_compilation_profiler",
         ":device_compiler",
+        ":pjrt_device_compiler_client",
         ":variable_info",
         ":variable_info_util",
         ":xla_device_no_jit_rewrite_registration",
         ":xla_launch_util",
         "//tensorflow/compiler/tf2xla:xla_compiler",
+        "//tensorflow/compiler/tf2xla:xla_resource",
         "//tensorflow/core:framework",
-        "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/common_runtime:core_cpu_internal",
         "//tensorflow/core/common_runtime/eager:tensor_handle",
         "//tensorflow/core/platform:refcount",
+        "//tensorflow/core/platform:statusor",
+        "//tensorflow/core/tpu:tpu_defs",
         "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",
+        "@xla//xla:shape_util",
         "@xla//xla:status_macros",
-        "@xla//xla/client:executable_build_options",
-        "@xla//xla/client:local_client",
+        "@xla//xla/hlo/ir:hlo",
         "@xla//xla/hlo/translate:portable_api",
+        "@xla//xla/pjrt:compiled_memory_stats",
+        "@xla//xla/pjrt:pjrt_client",
         "@xla//xla/service:hlo_graph_dumper",
+        "@xla//xla/service:hlo_module_config",
+        "@xla//xla/service:hlo_proto_util",
         "@xla//xla/stream_executor:platform",
         "@xla//xla/stream_executor/host:host_platform_id",
+        "@xla//xla/tsl/platform:errors",
+        "@xla//xla/tsl/platform:statusor",
     ],
     alwayslink = 1,
 )
@@ -1984,9 +1976,12 @@ tf_cuda_cc_test(
     ],
     deps = [
         ":device_compilation_cluster_signature",
+        ":device_compilation_profiler",
         ":device_compiler",
         ":device_compiler_client",
+        ":device_executable_persistor",
         ":tf_graph_to_hlo_compiler_proto_cc",
+        ":xla_activity_listener",
         ":xla_compile_util",
         ":xla_device_compiler_client",
         ":xla_gpu_device",
@@ -2004,14 +1999,17 @@ tf_cuda_cc_test(
         "//tensorflow/core/kernels:ops_testutil",
         "//tensorflow/core/platform:env",
         "//tensorflow/core/platform:errors",
+        "//tensorflow/core/platform:path",
         "//tensorflow/core/platform:refcount",
         "//tensorflow/core/platform:status",
         "//tensorflow/core/platform:status_matchers",
         "//tensorflow/core/platform:statusor",
         "//tensorflow/core/platform:test",
         "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/base",
         "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/time",
@@ -2022,6 +2020,7 @@ tf_cuda_cc_test(
         "@xla//xla/hlo/builder:xla_computation",
         "@xla//xla/stream_executor:platform_manager",
         "@xla//xla/tsl/lib/strings:proto_serialization",
+        "@xla//xla/tsl/platform:errors",
         "@xla//xla/tsl/platform:statusor",
     ],
 )
@@ -2075,8 +2074,10 @@ tf_cuda_cc_test(
         "@xla//xla/client:client_library",
         "@xla//xla/client:local_client",
         "@xla//xla/pjrt:pjrt_client",
-        "@xla//xla/stream_executor:device_memory_allocator",
+        "@xla//xla/service/cpu:cpu_transfer_manager",
+        "@xla//xla/stream_executor:device_address_allocator",
         "@xla//xla/stream_executor:platform",
+        "@xla//xla/stream_executor:platform_manager",
         "@xla//xla/stream_executor/host:host_platform_id",
         "@xla//xla/tsl/platform:statusor",
     ],