| 607 | } |
| 608 | |
| 609 | StatusOr<std::unique_ptr<PyTpuBuffer>> PyTpuExecutable::Execute( |
| 610 | absl::Span<PyTpuBuffer* const> argument_handles) { |
| 611 | if (num_replicas() != 1) { |
| 612 | return InvalidArgument( |
| 613 | "Attempted to execute computation with %d replicas using Execute()", |
| 614 | num_replicas()); |
| 615 | } |
| 616 | if (num_partitions() != 1) { |
| 617 | return InvalidArgument( |
| 618 | "Attempted to execute computation with %d partitions using Execute()", |
| 619 | num_partitions()); |
| 620 | } |
| 621 | |
| 622 | std::vector<PyTpuBuffer*> all_core_arguments(argument_handles.begin(), |
| 623 | argument_handles.end()); |
| 624 | |
| 625 | ExecuteResult result = |
| 626 | ExecuteHelper(absl::MakeSpan(&all_core_arguments, 1), argument_handles, |
| 627 | /*replica=*/0, /*partition=*/0, RunId()); |
| 628 | |
| 629 | Status status = WaitForExecuteEvent(result.on_execute_finished.get()); |
| 630 | |
| 631 | if (!status.ok()) { |
| 632 | LOG(ERROR) << "Failed to execute program: " << status; |
| 633 | return status; |
| 634 | } |
| 635 | |
| 636 | return std::move(result.buffer); |
| 637 | } |
| 638 | |
| 639 | StatusOr<std::vector<std::unique_ptr<PyTpuBuffer>>> |
| 640 | PyTpuExecutable::ExecutePerReplica( |
nothing calls this directly
no test coverage detected