| 425 | } |
| 426 | |
| 427 | StatusOr<std::vector<std::unique_ptr<PyLocalBuffer>>> |
| 428 | PyLocalBuffer::DestructureTuple() { |
| 429 | tensorflow::profiler::TraceMe traceme("PyLocalBuffer::DestructureTuple"); |
| 430 | absl::MutexLock lock(&mu_); |
| 431 | if (!on_host_shape_.IsTuple()) { |
| 432 | return InvalidArgument( |
| 433 | "Attempted to destructure a PyLocalBuffer that did not have a tuple " |
| 434 | "shape; shape: %s", |
| 435 | ShapeUtil::HumanString(on_host_shape_)); |
| 436 | } |
| 437 | if (!device_buffer_) { |
| 438 | return InvalidArgument("Attempted to destructure a deleted buffer."); |
| 439 | } |
| 440 | int num_children = ShapeUtil::TupleElementCount(on_host_shape_); |
| 441 | std::vector<std::unique_ptr<PyLocalBuffer>> results; |
| 442 | results.reserve(num_children); |
| 443 | for (int64 i = 0; i < num_children; ++i) { |
| 444 | results.push_back(absl::make_unique<PyLocalBuffer>( |
| 445 | on_host_shape_.tuple_shapes(i), on_device_shape_.tuple_shapes(i), |
| 446 | device_buffer_->children().at(i), client_, device_)); |
| 447 | } |
| 448 | return results; |
| 449 | } |
| 450 | |
| 451 | StatusOr<std::unique_ptr<PyLocalBuffer>> PyLocalBuffer::CopyToDevice( |
| 452 | std::shared_ptr<Device> dst_device) { |