From c5dc28ec6584f74f0af31b67f1828e2e761a3615 Mon Sep 17 00:00:00 2001
From: Sergei Lebedev <slebedev@google.com>
Date: Mon, 27 Jan 2025 06:02:33 -0800
Subject: [PATCH] [pjrt] Added a PjRtMemorySpace* overload for
 CreateUninitializedBuffer

PiperOrigin-RevId: 720139737
---
 xla/pjrt/pjrt_client.h                 | 10 +++++++++-
 xla/pjrt/pjrt_stream_executor_client.h |  1 +
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/xla/pjrt/pjrt_client.h b/xla/pjrt/pjrt_client.h
index 25f0e2f1600733..4c903b20c6cb94 100644
--- a/xla/pjrt/pjrt_client.h
+++ b/xla/pjrt/pjrt_client.h
@@ -641,9 +641,17 @@ class PjRtClient {
   }
 
   // Creates a buffer on the device without initializing or copying any data.
+  ABSL_DEPRECATED("Use CreateUninitializedBuffer(Shape, PjRtMemorySpace*)")
   virtual absl::StatusOr<std::unique_ptr<PjRtBuffer>> CreateUninitializedBuffer(
       const Shape& shape, PjRtDevice* device) {
-    return Unimplemented("CreateUnitializedBuffer is not supported.");
+    return Unimplemented("CreateUninitializedBuffer is not supported.");
+  }
+
+  // Creates a buffer in the given memory space without initializing or copying
+  // any data.
+  virtual absl::StatusOr<std::unique_ptr<PjRtBuffer>> CreateUninitializedBuffer(
+      const Shape& shape, PjRtMemorySpace* memory_space) {
+    return Unimplemented("CreateUninitializedBuffer is not supported.");
   }
 
   // Creates buffer in the given memory space that carries an error future
diff --git a/xla/pjrt/pjrt_stream_executor_client.h b/xla/pjrt/pjrt_stream_executor_client.h
index 0a4a7c430ab9cb..101e523a507dcb 100644
--- a/xla/pjrt/pjrt_stream_executor_client.h
+++ b/xla/pjrt/pjrt_stream_executor_client.h
@@ -333,6 +333,7 @@ class PjRtStreamExecutorClient : public PjRtClient {
   // initialized the data.
   absl::StatusOr<std::unique_ptr<PjRtBuffer>> CreateUninitializedBuffer(
       const Shape& shape, PjRtDevice* device) override;
+  using PjRtClient::CreateUninitializedBuffer;
   absl::StatusOr<std::unique_ptr<PjRtBuffer>> CreateUninitializedBuffer(
       const Shape& shape, PjRtDevice* device,
       std::shared_ptr<BufferSequencingEvent> definition_event);