[pjrt] Added a PjRtMemorySpace* overload for CreateUninitializedBuffer

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

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);