| 563 | } |
| 564 | |
| 565 | void XRTTupleAllocation::InitializeFromShapedBuffer( |
| 566 | const xla::ShapedBuffer& shaped_buffer, |
| 567 | se::DeviceMemoryAllocator* allocator, int device_ordinal) { |
| 568 | for (auto& index_buffer : buffers_) { |
| 569 | if (index_buffer.second != nullptr) { |
| 570 | index_buffer.second->Unref(); |
| 571 | } |
| 572 | // Make a reference-counted version of the allocated buffer. |
| 573 | index_buffer.second = new XRTBufferAllocation( |
| 574 | shaped_buffer.buffer(index_buffer.first), device_ordinal, allocator); |
| 575 | } |
| 576 | } |
| 577 | |
| 578 | xla::StatusOr<xla::ShapedBuffer> XRTTupleAllocation::ToShapedBuffer() { |
| 579 | xla::ShapedBuffer shaped_buffer(on_host_shape(), on_device_shape(), |
no test coverage detected