| 639 | } |
| 640 | |
| 641 | xla::StatusOr<xla::ShapeTree<xla::MaybeOwningDeviceMemory>> |
| 642 | XRTTupleAllocation::ToDeviceMemoryTree( |
| 643 | const std::function<bool(const xla::ShapeIndex&)>& release_checker) { |
| 644 | xla::ShapeTree<xla::MaybeOwningDeviceMemory> shaped_tree(on_device_shape()); |
| 645 | for (const auto& index_buffer : buffers_) { |
| 646 | if (index_buffer.second == nullptr || |
| 647 | index_buffer.second->allocation().is_null()) { |
| 648 | return errors::InvalidArgument("Literal buffer at index ", |
| 649 | index_buffer.first.ToString(), |
| 650 | " has been released"); |
| 651 | } |
| 652 | if (!release_checker(index_buffer.first)) { |
| 653 | *shaped_tree.mutable_element(index_buffer.first) = |
| 654 | index_buffer.second->allocation(); |
| 655 | } else { |
| 656 | // We keep the ownership of the device memory here. |
| 657 | *shaped_tree.mutable_element(index_buffer.first) = se::OwningDeviceMemory( |
| 658 | index_buffer.second->allocation(), device_ordinal_, allocator_); |
| 659 | } |
| 660 | } |
| 661 | return std::move(shaped_tree); |
| 662 | } |
| 663 | |
| 664 | } // namespace tensorflow |
nothing calls this directly
no test coverage detected