Automated Code Change

PiperOrigin-RevId: 707810663
openxla · Dec 19, 2024 · 37e7fb0 · 37e7fb0
1 parent 926ef6a
commit 37e7fb0
Show file tree

Hide file tree

Showing 7 changed files with 14 additions and 0 deletions.
diff --git a/xla/python/BUILD b/xla/python/BUILD
@@ -605,6 +605,7 @@ cc_library(
     srcs = ["inspect_sharding.cc"],
     hdrs = ["inspect_sharding.h"],
     deps = [
+        "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/service:custom_call_sharding_helper",
         "//xla/service/spmd:spmd_partitioner",
@@ -703,6 +704,7 @@ cc_library(
         "//xla/hlo/builder/lib:sorting",
         "//xla/hlo/builder/lib:svd",
         "//xla/pjrt:status_casters",
+        "//xla/service:hlo_proto_cc",
     ],
 )
 
@@ -1021,6 +1023,9 @@ cc_library(
         "@tsl//tsl/profiler/lib:profiler_interface",
         "@tsl//tsl/profiler/lib:profiler_session",
         "@tsl//tsl/profiler/lib:traceme",
+        "@tsl//tsl/profiler/protobuf:profiled_instructions_proto_cc",
+        "@tsl//tsl/profiler/protobuf:profiler_options_proto_cc",
+        "@tsl//tsl/profiler/protobuf:xplane_proto_cc",
     ],
 )
 

diff --git a/xla/python/inspect_sharding.cc b/xla/python/inspect_sharding.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "xla/python/inspect_sharding.h"
 
+#include <cstring>
 #include <limits>
 #include <memory>
 #include <optional>
@@ -24,6 +25,7 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "xla/service/custom_call_sharding_helper.h"
 #include "xla/service/spmd/spmd_partitioner_util.h"
+#include "xla/xla_data.pb.h"
 
 namespace jax {
 

diff --git a/xla/python/inspect_sharding.h b/xla/python/inspect_sharding.h
@@ -16,6 +16,7 @@ limitations under the License.
 #ifndef XLA_PYTHON_INSPECT_SHARDING_H_
 #define XLA_PYTHON_INSPECT_SHARDING_H_
 
+#include <cstddef>
 #include <optional>
 #include <string>
 

diff --git a/xla/python/jax_jit.h b/xla/python/jax_jit.h
@@ -19,6 +19,7 @@ limitations under the License.
 #include <Python.h>
 
 #include <cstddef>
+#include <memory>
 #include <optional>
 #include <stdexcept>
 #include <string>

diff --git a/xla/python/ops.cc b/xla/python/ops.cc
@@ -46,6 +46,7 @@ limitations under the License.
 #include "xla/python/nb_absl_span.h"  // IWYU pragma: keep
 #include "xla/python/nb_helpers.h"
 #include "xla/python/types.h"
+#include "xla/service/hlo.pb.h"
 #include "xla/xla_data.pb.h"
 
 namespace nb = nanobind;

diff --git a/xla/python/pprof_profile_builder.cc b/xla/python/pprof_profile_builder.cc
@@ -26,6 +26,7 @@ limitations under the License.
 #include "xla/util.h"
 #include "tsl/platform/logging.h"
 #include "tsl/platform/protobuf.h"
+#include "tsl/profiler/protobuf/profile.pb.h"
 
 namespace xla {
 

diff --git a/xla/python/profiler.cc b/xla/python/profiler.cc
@@ -42,6 +42,9 @@ limitations under the License.
 #include "tsl/platform/protobuf.h"  // IWYU pragma: keep
 #include "tsl/profiler/lib/profiler_session.h"
 #include "tsl/profiler/lib/traceme.h"
+#include "tsl/profiler/protobuf/profiled_instructions.pb.h"
+#include "tsl/profiler/protobuf/profiler_options.pb.h"
+#include "tsl/profiler/protobuf/xplane.pb.h"
 
 namespace xla {