Skip to content

Commit

Permalink
[pjrt] Added a PjRtMemorySpace* overload for CreateUninitializedBuffer
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 720578965
  • Loading branch information
superbobry authored and Google-ML-Automation committed Jan 28, 2025
1 parent b215246 commit f133f0b
Showing 2 changed files with 10 additions and 1 deletion.
10 changes: 9 additions & 1 deletion xla/pjrt/pjrt_client.h
Original file line number Diff line number Diff line change
@@ -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
1 change: 1 addition & 0 deletions xla/pjrt/pjrt_stream_executor_client.h
Original file line number Diff line number Diff line change
@@ -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);

0 comments on commit f133f0b

Please sign in to comment.