tf-nightly-cpu 2.20.0.dev20250220__cp310-cp310-win_amd64.whl → 2.20.0.dev20250222__cp310-cp310-win_amd64.whl

tensorflow/include/xla/mlir_hlo/stablehlo_ext/transforms/passes.h.inc

@@ -3,6 +3,7 @@
 #ifdef GEN_PASS_DECL
 // Generate declarations for all passes.
 #define GEN_PASS_DECL_CHLORECOMPOSEOPSPASS
+#define GEN_PASS_DECL_STABLEHLOADDQDQAFTERCONVPASS
 #define GEN_PASS_DECL_STABLEHLOCANONICALIZEDYNAMISMPASS
 #define GEN_PASS_DECL_STABLEHLOFLATTENENTRYFUNCTIONTUPLESPASS
 #define GEN_PASS_DECL_STABLEHLOFLATTENTUPLEPASS
@@ -87,6 +88,82 @@ std::unique_ptr<::mlir::Pass> createChloRecomposeOpsPass() {
 #undef GEN_PASS_DEF_CHLORECOMPOSEOPSPASS
 #endif // GEN_PASS_DEF_CHLORECOMPOSEOPSPASS
 
+//===----------------------------------------------------------------------===//
+// StablehloAddQDQAfterConvPass
+//===----------------------------------------------------------------------===//
+#ifdef GEN_PASS_DECL_STABLEHLOADDQDQAFTERCONVPASS
+std::unique_ptr<::mlir::Pass> createStablehloAddQDQAfterConvPass();
+#undef GEN_PASS_DECL_STABLEHLOADDQDQAFTERCONVPASS
+#endif // GEN_PASS_DECL_STABLEHLOADDQDQAFTERCONVPASS
+#ifdef GEN_PASS_DEF_STABLEHLOADDQDQAFTERCONVPASS
+
+namespace impl {
+  std::unique_ptr<::mlir::Pass> createStablehloAddQDQAfterConvPass();
+} // namespace impl
+namespace impl {
+
+template <typename DerivedT>
+class StablehloAddQDQAfterConvPassBase : public ::mlir::OperationPass<ModuleOp> {
+public:
+  using Base = StablehloAddQDQAfterConvPassBase;
+
+  StablehloAddQDQAfterConvPassBase() : ::mlir::OperationPass<ModuleOp>(::mlir::TypeID::get<DerivedT>()) {}
+  StablehloAddQDQAfterConvPassBase(const StablehloAddQDQAfterConvPassBase &other) : ::mlir::OperationPass<ModuleOp>(other) {}
+  StablehloAddQDQAfterConvPassBase& operator=(const StablehloAddQDQAfterConvPassBase &) = delete;
+  StablehloAddQDQAfterConvPassBase(StablehloAddQDQAfterConvPassBase &&) = delete;
+  StablehloAddQDQAfterConvPassBase& operator=(StablehloAddQDQAfterConvPassBase &&) = delete;
+  ~StablehloAddQDQAfterConvPassBase() = default;
+
+  /// Returns the command-line argument attached to this pass.
+  static constexpr ::llvm::StringLiteral getArgumentName() {
+    return ::llvm::StringLiteral("stablehlo-ext-add-qdq-after-conv");
+  }
+  ::llvm::StringRef getArgument() const override { return "stablehlo-ext-add-qdq-after-conv"; }
+
+  ::llvm::StringRef getDescription() const override { return "Add quant and dequant ops after convolution op."; }
+
+  /// Returns the derived pass name.
+  static constexpr ::llvm::StringLiteral getPassName() {
+    return ::llvm::StringLiteral("StablehloAddQDQAfterConvPass");
+  }
+  ::llvm::StringRef getName() const override { return "StablehloAddQDQAfterConvPass"; }
+
+  /// Support isa/dyn_cast functionality for the derived pass class.
+  static bool classof(const ::mlir::Pass *pass) {
+    return pass->getTypeID() == ::mlir::TypeID::get<DerivedT>();
+  }
+
+  /// A clone method to create a copy of this pass.
+  std::unique_ptr<::mlir::Pass> clonePass() const override {
+    return std::make_unique<DerivedT>(*static_cast<const DerivedT *>(this));
+  }
+
+  /// Return the dialect that must be loaded in the context before this pass.
+  void getDependentDialects(::mlir::DialectRegistry &registry) const override {
+    registry.insert<mlir::quant::QuantDialect>();
+    registry.insert<stablehlo::StablehloDialect>();
+  }
+
+  /// Explicitly declare the TypeID for this class. We declare an explicit private
+  /// instantiation because Pass classes should only be visible by the current
+  /// library.
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(StablehloAddQDQAfterConvPassBase<DerivedT>)
+
+protected:
+private:
+
+  friend std::unique_ptr<::mlir::Pass> createStablehloAddQDQAfterConvPass() {
+    return std::make_unique<DerivedT>();
+  }
+};
+} // namespace impl
+
+std::unique_ptr<::mlir::Pass> createStablehloAddQDQAfterConvPass() {
+  return impl::createStablehloAddQDQAfterConvPass();
+}
+#undef GEN_PASS_DEF_STABLEHLOADDQDQAFTERCONVPASS
+#endif // GEN_PASS_DEF_STABLEHLOADDQDQAFTERCONVPASS
+
 //===----------------------------------------------------------------------===//
 // StablehloCanonicalizeDynamismPass
 //===----------------------------------------------------------------------===//
@@ -360,9 +437,9 @@ public:
 
   /// Returns the command-line argument attached to this pass.
   static constexpr ::llvm::StringLiteral getArgumentName() {
-    return ::llvm::StringLiteral("legalize-quant-composite");
+    return ::llvm::StringLiteral("stablehlo-ext-legalize-quant-composite");
   }
-  ::llvm::StringRef getArgument() const override { return "legalize-quant-composite"; }
+  ::llvm::StringRef getArgument() const override { return "stablehlo-ext-legalize-quant-composite"; }
 
   ::llvm::StringRef getDescription() const override { return "Lowers the quantization related composites op to native quantized ops."; }
 
@@ -576,6 +653,23 @@ inline void registerChloRecomposeOpsPassPass() {
   });
 }
 
+//===----------------------------------------------------------------------===//
+// StablehloAddQDQAfterConvPass Registration
+//===----------------------------------------------------------------------===//
+
+inline void registerStablehloAddQDQAfterConvPass() {
+  ::mlir::registerPass([]() -> std::unique_ptr<::mlir::Pass> {
+    return createStablehloAddQDQAfterConvPass();
+  });
+}
+
+// Old registration code, kept for temporary backwards compatibility.
+inline void registerStablehloAddQDQAfterConvPassPass() {
+  ::mlir::registerPass([]() -> std::unique_ptr<::mlir::Pass> {
+    return createStablehloAddQDQAfterConvPass();
+  });
+}
+
 //===----------------------------------------------------------------------===//
 // StablehloCanonicalizeDynamismPass Registration
 //===----------------------------------------------------------------------===//
@@ -684,6 +778,7 @@ inline void registerStablehloRefineShapesPassPass() {
 
 inline void registerPasses() {
   registerChloRecomposeOpsPass();
+  registerStablehloAddQDQAfterConvPass();
   registerStablehloCanonicalizeDynamismPass();
   registerStablehloFlattenEntryFunctionTuplesPass();
   registerStablehloFlattenTuplePass();
@@ -745,6 +840,56 @@ public:
 protected:
 };
 
+template <typename DerivedT>
+class StablehloAddQDQAfterConvPassBase : public ::mlir::OperationPass<ModuleOp> {
+public:
+  using Base = StablehloAddQDQAfterConvPassBase;
+
+  StablehloAddQDQAfterConvPassBase() : ::mlir::OperationPass<ModuleOp>(::mlir::TypeID::get<DerivedT>()) {}
+  StablehloAddQDQAfterConvPassBase(const StablehloAddQDQAfterConvPassBase &other) : ::mlir::OperationPass<ModuleOp>(other) {}
+  StablehloAddQDQAfterConvPassBase& operator=(const StablehloAddQDQAfterConvPassBase &) = delete;
+  StablehloAddQDQAfterConvPassBase(StablehloAddQDQAfterConvPassBase &&) = delete;
+  StablehloAddQDQAfterConvPassBase& operator=(StablehloAddQDQAfterConvPassBase &&) = delete;
+  ~StablehloAddQDQAfterConvPassBase() = default;
+
+  /// Returns the command-line argument attached to this pass.
+  static constexpr ::llvm::StringLiteral getArgumentName() {
+    return ::llvm::StringLiteral("stablehlo-ext-add-qdq-after-conv");
+  }
+  ::llvm::StringRef getArgument() const override { return "stablehlo-ext-add-qdq-after-conv"; }
+
+  ::llvm::StringRef getDescription() const override { return "Add quant and dequant ops after convolution op."; }
+
+  /// Returns the derived pass name.
+  static constexpr ::llvm::StringLiteral getPassName() {
+    return ::llvm::StringLiteral("StablehloAddQDQAfterConvPass");
+  }
+  ::llvm::StringRef getName() const override { return "StablehloAddQDQAfterConvPass"; }
+
+  /// Support isa/dyn_cast functionality for the derived pass class.
+  static bool classof(const ::mlir::Pass *pass) {
+    return pass->getTypeID() == ::mlir::TypeID::get<DerivedT>();
+  }
+
+  /// A clone method to create a copy of this pass.
+  std::unique_ptr<::mlir::Pass> clonePass() const override {
+    return std::make_unique<DerivedT>(*static_cast<const DerivedT *>(this));
+  }
+
+  /// Register the dialects that must be loaded in the context before this pass.
+  void getDependentDialects(::mlir::DialectRegistry &registry) const override {
+    registry.insert<mlir::quant::QuantDialect>();
+    registry.insert<stablehlo::StablehloDialect>();
+  }
+
+  /// Explicitly declare the TypeID for this class. We declare an explicit private
+  /// instantiation because Pass classes should only be visible by the current
+  /// library.
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(StablehloAddQDQAfterConvPassBase<DerivedT>)
+
+protected:
+};
+
 template <typename DerivedT>
 class StablehloCanonicalizeDynamismPassBase : public ::mlir::OperationPass<func::FuncOp> {
 public:
@@ -907,9 +1052,9 @@ public:
 
   /// Returns the command-line argument attached to this pass.
   static constexpr ::llvm::StringLiteral getArgumentName() {
-    return ::llvm::StringLiteral("legalize-quant-composite");
+    return ::llvm::StringLiteral("stablehlo-ext-legalize-quant-composite");
   }
-  ::llvm::StringRef getArgument() const override { return "legalize-quant-composite"; }
+  ::llvm::StringRef getArgument() const override { return "stablehlo-ext-legalize-quant-composite"; }
 
   ::llvm::StringRef getDescription() const override { return "Lowers the quantization related composites op to native quantized ops."; }
 

tensorflow/include/xla/pjrt/distributed/client.h

@@ -145,6 +145,11 @@ class DistributedRuntimeClient {
       std::string barrier_id, absl::Duration timeout,
       std::optional<absl::Span<const int32_t>> nodes) = 0;
 
+  // Returns the subset of live nodes. See CoordinationService.GetAliveTasks for
+  // detailed semantics.
+  virtual absl::StatusOr<std::vector<int32_t>> GetLiveNodes(
+      absl::Span<const int32_t> nodes) = 0;
+
   // Returns pointer to coordination service agent, or InternalError if the
   // client does not use coordination service.
   virtual absl::StatusOr<tsl::CoordinationServiceAgent*>

tensorflow/include/xla/pjrt/gpu/se_gpu_pjrt_client.h

@@ -37,6 +37,7 @@ limitations under the License.
 #include "xla/pjrt/distributed/key_value_store_interface.h"
 #include "xla/pjrt/gpu/gpu_topology.h"
 #include "xla/pjrt/gpu/gpu_topology.pb.h"
+#include "xla/pjrt/gpu/se_gpu_topology_description.h"
 #include "xla/pjrt/local_device_state.h"
 #include "xla/pjrt/pjrt_client.h"
 #include "xla/pjrt/pjrt_compiler.h"
@@ -57,98 +58,6 @@ using DeviceTopologyPair =
     std::pair<std::vector<std::unique_ptr<PjRtStreamExecutorDevice>>,
               GpuTopologyProto>;
 
-class StreamExecutorGpuTopologyDescription : public PjRtTopologyDescription {
- public:
-  StreamExecutorGpuTopologyDescription(
-      const PjRtPlatformId platform_id, const absl::string_view platform_name,
-      std::shared_ptr<const GpuTopology> gpu_topology,
-      const absl::flat_hash_map<std::string, PjRtDeviceAttribute>& attributes =
-          {},
-      std::optional<stream_executor::GpuTargetConfigProto> target_config =
-          std::nullopt)
-      : platform_id_(platform_id),
-        platform_name_(platform_name),
-        gpu_topology_(std::move(gpu_topology)),
-        attributes_(attributes),
-        target_config_(std::move(target_config)) {}
-
-  bool operator==(const StreamExecutorGpuTopologyDescription& other) const {
-    return this->platform_id() == other.platform_id() &&
-           this->platform_name() == other.platform_name() &&
-           this->platform_version() == other.platform_version() &&
-           this->gpu_topology() == other.gpu_topology();
-  }
-
-  PjRtPlatformId platform_id() const override { return platform_id_; }
-
-  absl::string_view platform_name() const override { return platform_name_; }
-
-  absl::string_view platform_version() const override {
-    return gpu_topology_->platform_version();
-  }
-
-  std::vector<std::unique_ptr<const PjRtDeviceDescription>> DeviceDescriptions()
-      const override {
-    std::vector<std::unique_ptr<const PjRtDeviceDescription>> devices;
-    devices.reserve(gpu_topology_->number_of_devices());
-    for (const int device_id : gpu_topology_->device_ids()) {
-      devices.push_back(std::make_unique<PjRtStreamExecutorDeviceDescription>(
-          device_id, std::string(platform_version())));
-    }
-    return devices;
-  }
-
-  const GpuTopology& gpu_topology() const { return *gpu_topology_; }
-  const GpuTopology* gpu_topology_ptr() const { return gpu_topology_.get(); }
-
-  // No subslice is supported.
-  bool is_subslice_topology() const override { return false; }
-
-  absl::StatusOr<int> ProcessCount() const override {
-    return gpu_topology_->number_of_hosts();
-  }
-
-  absl::StatusOr<int> CoreCountOfDefaultType() const override {
-    return gpu_topology_->number_of_devices();
-  }
-
-  absl::StatusOr<int> LogicalDeviceCountOfDefaultType() const override {
-    return gpu_topology_->number_of_devices();
-  }
-
-  absl::StatusOr<int> CoreCountOfDefaultTypePerProcess() const override {
-    return gpu_topology_->number_of_devices();
-  }
-
-  absl::StatusOr<int> CoreCountOfDefaultTypePerChip() const override {
-    return 1;
-  }
-
-  absl::StatusOr<std::string> Serialize() const override;
-
-  const std::optional<stream_executor::GpuTargetConfigProto>& target_config()
-      const {
-    return target_config_;
-  }
-
-  // Returns vendor specific attributes about the topology.
-  const absl::flat_hash_map<std::string, PjRtDeviceAttribute>& Attributes()
-      const override {
-    return attributes_;
-  }
-
-  absl::StatusOr<Layout> GetDefaultLayout(
-      PrimitiveType element_type,
-      absl::Span<const int64_t> dims) const override;
-
- private:
-  const PjRtPlatformId platform_id_;
-  const std::string platform_name_;
-  std::shared_ptr<const GpuTopology> gpu_topology_;
-  absl::flat_hash_map<std::string, xla::PjRtDeviceAttribute> attributes_;
-  std::optional<stream_executor::GpuTargetConfigProto> target_config_;
-};
-
 class StreamExecutorGpuDevice : public PjRtStreamExecutorDevice {
  public:
   StreamExecutorGpuDevice(int id,

tensorflow/include/xla/pjrt/gpu/se_gpu_topology_description.h

@@ -0,0 +1,126 @@
+/* Copyright 2025 The OpenXLA Authors.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef XLA_PJRT_GPU_SE_GPU_TOPOLOGY_DESCRIPTION_H_
+#define XLA_PJRT_GPU_SE_GPU_TOPOLOGY_DESCRIPTION_H_
+
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "xla/pjrt/gpu/gpu_topology.h"
+#include "xla/pjrt/pjrt_compiler.h"
+#include "xla/pjrt/pjrt_device_description.h"
+#include "xla/pjrt/pjrt_stream_executor_device_description.h"
+
+namespace xla {
+
+class StreamExecutorGpuTopologyDescription : public PjRtTopologyDescription {
+ public:
+  StreamExecutorGpuTopologyDescription(
+      const PjRtPlatformId platform_id, const absl::string_view platform_name,
+      std::shared_ptr<const GpuTopology> gpu_topology,
+      const absl::flat_hash_map<std::string, PjRtDeviceAttribute>& attributes =
+          {},
+      std::optional<stream_executor::GpuTargetConfigProto> target_config =
+          std::nullopt)
+      : platform_id_(platform_id),
+        platform_name_(platform_name),
+        gpu_topology_(std::move(gpu_topology)),
+        attributes_(attributes),
+        target_config_(std::move(target_config)) {}
+
+  bool operator==(const StreamExecutorGpuTopologyDescription& other) const {
+    return this->platform_id() == other.platform_id() &&
+           this->platform_name() == other.platform_name() &&
+           this->platform_version() == other.platform_version() &&
+           this->gpu_topology() == other.gpu_topology();
+  }
+
+  PjRtPlatformId platform_id() const override { return platform_id_; }
+
+  absl::string_view platform_name() const override { return platform_name_; }
+
+  absl::string_view platform_version() const override {
+    return gpu_topology_->platform_version();
+  }
+
+  std::vector<std::unique_ptr<const PjRtDeviceDescription>> DeviceDescriptions()
+      const override {
+    std::vector<std::unique_ptr<const PjRtDeviceDescription>> devices;
+    devices.reserve(gpu_topology_->number_of_devices());
+    for (const int device_id : gpu_topology_->device_ids()) {
+      devices.push_back(std::make_unique<PjRtStreamExecutorDeviceDescription>(
+          device_id, std::string(platform_version())));
+    }
+    return devices;
+  }
+
+  const GpuTopology& gpu_topology() const { return *gpu_topology_; }
+  const GpuTopology* gpu_topology_ptr() const { return gpu_topology_.get(); }
+
+  // No subslice is supported.
+  bool is_subslice_topology() const override { return false; }
+
+  absl::StatusOr<int> ProcessCount() const override {
+    return gpu_topology_->number_of_hosts();
+  }
+
+  absl::StatusOr<int> CoreCountOfDefaultType() const override {
+    return gpu_topology_->number_of_devices();
+  }
+
+  absl::StatusOr<int> LogicalDeviceCountOfDefaultType() const override {
+    return gpu_topology_->number_of_devices();
+  }
+
+  absl::StatusOr<int> CoreCountOfDefaultTypePerProcess() const override {
+    return gpu_topology_->number_of_devices();
+  }
+
+  absl::StatusOr<int> CoreCountOfDefaultTypePerChip() const override {
+    return 1;
+  }
+
+  absl::StatusOr<std::string> Serialize() const override;
+
+  const std::optional<stream_executor::GpuTargetConfigProto>& target_config()
+      const {
+    return target_config_;
+  }
+
+  // Returns vendor specific attributes about the topology.
+  const absl::flat_hash_map<std::string, PjRtDeviceAttribute>& Attributes()
+      const override {
+    return attributes_;
+  }
+
+  absl::StatusOr<Layout> GetDefaultLayout(
+      PrimitiveType element_type,
+      absl::Span<const int64_t> dims) const override;
+
+ private:
+  const PjRtPlatformId platform_id_;
+  const std::string platform_name_;
+  std::shared_ptr<const GpuTopology> gpu_topology_;
+  absl::flat_hash_map<std::string, xla::PjRtDeviceAttribute> attributes_;
+  std::optional<stream_executor::GpuTargetConfigProto> target_config_;
+};
+}  // namespace xla
+
+#endif  // XLA_PJRT_GPU_SE_GPU_TOPOLOGY_DESCRIPTION_H_

tensorflow/include/xla/pjrt/pjrt_stream_executor_client.h

@@ -52,9 +52,9 @@ limitations under the License.
 #include "xla/pjrt/pjrt_client.h"
 #include "xla/pjrt/pjrt_common.h"
 #include "xla/pjrt/pjrt_compiler.h"
-#include "xla/pjrt/pjrt_device_description.h"
 #include "xla/pjrt/pjrt_executable.h"
 #include "xla/pjrt/pjrt_future.h"
+#include "xla/pjrt/pjrt_stream_executor_device_description.h"
 #include "xla/pjrt/tracked_device_buffer.h"
 #include "xla/pjrt/transpose.h"
 #include "xla/pjrt/utils.h"
@@ -77,54 +77,6 @@ limitations under the License.
 
 namespace xla {
 
-class PjRtStreamExecutorDeviceDescription : public PjRtDeviceDescription {
- public:
-  explicit PjRtStreamExecutorDeviceDescription(int id, std::string device_kind,
-                                               int process_index = 0)
-      : id_(id),
-        process_index_(process_index),
-        device_kind_(std::move(device_kind)) {}
-
-  int id() const override { return id_; }
-
-  int process_index() const override { return process_index_; }
-
-  absl::string_view device_kind() const override { return device_kind_; }
-
-  absl::string_view ToString() const override { return to_string_; }
-
-  absl::string_view DebugString() const override { return debug_string_; }
-
-  absl::Span<int const> coords() const { return absl::MakeSpan(coords_); }
-
-  const absl::flat_hash_map<std::string, PjRtDeviceAttribute>& Attributes()
-      const override {
-    return attributes_;
-  }
-
-  void SetAttributes(
-      absl::flat_hash_map<std::string, PjRtDeviceAttribute> attributes) {
-    attributes_ = std::move(attributes);
-  }
-
-  void SetDebugString(std::string debug_string) {
-    debug_string_ = std::move(debug_string);
-  }
-
-  void SetToString(std::string to_string) { to_string_ = std::move(to_string); }
-
-  void SetCoords(std::array<int, 1> coords) { coords_ = coords; }
-
- private:
-  const int id_;
-  const int process_index_;
-  const std::string device_kind_;
-  std::string debug_string_ = "<unknown SE device>";
-  std::string to_string_ = "<unknown SE device>";
-  absl::flat_hash_map<std::string, PjRtDeviceAttribute> attributes_;
-  std::array<int, 1> coords_;
-};
-
 class PjRtStreamExecutorDevice : public PjRtDevice {
  public:
   explicit PjRtStreamExecutorDevice(

tensorflow/include/xla/pjrt/pjrt_stream_executor_device_description.h

@@ -0,0 +1,75 @@
+/* Copyright 2025 The OpenXLA Authors.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef XLA_PJRT_PJRT_STREAM_EXECUTOR_DEVICE_DESCRIPTION_H_
+#define XLA_PJRT_PJRT_STREAM_EXECUTOR_DEVICE_DESCRIPTION_H_
+
+#include <array>
+#include <string>
+#include <utility>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "xla/pjrt/pjrt_device_description.h"
+
+namespace xla {
+
+class PjRtStreamExecutorDeviceDescription : public PjRtDeviceDescription {
+ public:
+  explicit PjRtStreamExecutorDeviceDescription(int id, std::string device_kind,
+                                               int process_index = 0)
+      : id_(id),
+        process_index_(process_index),
+        device_kind_(std::move(device_kind)) {}
+
+  int id() const override { return id_; }
+
+  int process_index() const override { return process_index_; }
+
+  absl::string_view device_kind() const override { return device_kind_; }
+
+  absl::string_view ToString() const override { return to_string_; }
+
+  absl::string_view DebugString() const override { return debug_string_; }
+
+  absl::Span<int const> coords() const { return absl::MakeSpan(coords_); }
+
+  const absl::flat_hash_map<std::string, PjRtDeviceAttribute>& Attributes()
+      const override {
+    return attributes_;
+  }
+
+  void SetAttributes(
+      absl::flat_hash_map<std::string, PjRtDeviceAttribute> attributes) {
+    attributes_ = std::move(attributes);
+  }
+
+  void SetDebugString(std::string debug_string) {
+    debug_string_ = std::move(debug_string);
+  }
+
+  void SetToString(std::string to_string) { to_string_ = std::move(to_string); }
+
+  void SetCoords(std::array<int, 1> coords) { coords_ = coords; }
+
+ private:
+  const int id_;
+  const int process_index_;
+  const std::string device_kind_;
+  std::string debug_string_ = "<unknown SE device>";
+  std::string to_string_ = "<unknown SE device>";
+  absl::flat_hash_map<std::string, PjRtDeviceAttribute> attributes_;
+  std::array<int, 1> coords_;
+};
+}  // namespace xla
+
+#endif  // XLA_PJRT_PJRT_STREAM_EXECUTOR_DEVICE_DESCRIPTION_H_

tensorflow/include/xla/pjrt/plugin/xla_cpu/cpu_execute_options.h

@@ -0,0 +1,57 @@
+/* Copyright 2025 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef XLA_PJRT_PLUGIN_XLA_CPU_CPU_EXECUTE_OPTIONS_H_
+#define XLA_PJRT_PLUGIN_XLA_CPU_CPU_EXECUTE_OPTIONS_H_
+
+#include <optional>
+
+#include "xla/backends/cpu/collectives/cpu_collectives.h"
+#include "xla/pjrt/pjrt_executable.h"
+
+namespace xla {
+
+// ExecuteContext for XLA:CPU PjRtLoadedExecutable::Execute calls.
+class CpuExecuteContext : public ExecuteContext {
+ public:
+  ~CpuExecuteContext() override = default;
+
+  // If specified, override the process ID specified in
+  // `CpuClientOptions::process_id` for a particular call of
+  // PjRtLoadedExecutable::Execute.
+  //
+  // TODO(hyeontaek): Look for a collectives-agnostic way and combine this
+  // option with `ExecuteOptions::multi_slice_config`.
+  std::optional<int>& process_index() { return process_index_; }
+  std::optional<int> process_index() const { return process_index_; }
+
+  // If specified, override CPU collectives specified in
+  // `CpuClientOptions::collectives` for a particular call of
+  // PjRtLoadedExecutable::Execute. Must remain valid until the execution
+  // finishes.
+  //
+  // TODO(hyeontaek): Look for a collectives-agnostic way and combine this
+  // option with `ExecuteOptions::multi_slice_config`.
+  cpu::CpuCollectives*& collectives() { return collectives_; }
+  cpu::CpuCollectives* collectives() const { return collectives_; }
+
+ private:
+  std::optional<int> process_index_;
+  cpu::CpuCollectives* collectives_ = nullptr;
+};
+
+}  // namespace xla
+
+#endif  // XLA_PJRT_PLUGIN_XLA_CPU_CPU_EXECUTE_OPTIONS_H_

tensorflow/include/xla/pjrt/plugin/xla_cpu/cpu_topology.h

@@ -69,6 +69,10 @@ inline int UnpackCpuProcessIndex(PjRtGlobalDeviceId global_device_id) {
   return global_device_id.value() / kMaxCpuDevicesPerProcess;
 }
 
+inline int UnpackCpuLocalDeviceId(PjRtGlobalDeviceId global_device_id) {
+  return global_device_id.value() % kMaxCpuDevicesPerProcess;
+}
+
 }  // namespace xla
 
 #endif  // XLA_PJRT_PLUGIN_XLA_CPU_CPU_TOPOLOGY_H_

tensorflow/include/xla/service/constant_value.h

@@ -18,6 +18,7 @@ limitations under the License.
 
 #include <string>
 
+#include "absl/base/casts.h"
 #include "absl/status/statusor.h"
 #include "xla/literal.h"
 #include "xla/util.h"