diff --git a/xla/pjrt/pjrt_client.h b/xla/pjrt/pjrt_client.h
index 25f0e2f160073..4c903b20c6cb9 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 0a4a7c430ab9c..101e523a507dc 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);