static*/
| 841 | } |
| 842 | |
| 843 | /*static*/ StatusOr<std::unique_ptr<PyLocalExecutable>> |
| 844 | PyLocalExecutable::CompileForDevices( |
| 845 | const XlaComputation& computation, |
| 846 | absl::optional<std::vector<Shape>> argument_layouts, |
| 847 | const ExecutableBuildOptions* build_options, |
| 848 | std::shared_ptr<PyLocalClient> client, |
| 849 | const std::vector<std::vector<std::shared_ptr<Device>>>& |
| 850 | device_assignment) { |
| 851 | if (device_assignment.empty()) { |
| 852 | return InvalidArgument( |
| 853 | "Device assignment passed to Compile() must be non-empty."); |
| 854 | } |
| 855 | if (device_assignment[0].empty()) { |
| 856 | return InvalidArgument( |
| 857 | "Device assignment passed to Compile() must have a nonzero number of " |
| 858 | "partitions per replica; replica 0 had 0 partitions."); |
| 859 | } |
| 860 | DeviceAssignment xla_assignment(device_assignment.size(), |
| 861 | device_assignment[0].size()); |
| 862 | for (int replica = 0; replica < device_assignment.size(); ++replica) { |
| 863 | if (device_assignment[replica].size() != device_assignment[0].size()) { |
| 864 | return InvalidArgument( |
| 865 | "Device assignment passed to Compile() has different numbers of " |
| 866 | "partitions between replicas; %d partitions for replica %d versus %d " |
| 867 | "partitions for replica 0.", |
| 868 | device_assignment[replica].size(), replica, |
| 869 | device_assignment[0].size()); |
| 870 | } |
| 871 | for (int partition = 0; partition < device_assignment[replica].size(); |
| 872 | ++partition) { |
| 873 | if (device_assignment[0][0]->platform_name() != |
| 874 | device_assignment[replica][partition]->platform_name()) { |
| 875 | return InvalidArgument( |
| 876 | "Device assignment passed to Compile() must have devices of a " |
| 877 | "single kind, got %s for replica 0 partition 0 and %s for replica " |
| 878 | "%d partition %d.", |
| 879 | device_assignment[0][0]->platform_name(), |
| 880 | device_assignment[replica][partition]->platform_name(), replica, |
| 881 | partition); |
| 882 | } |
| 883 | xla_assignment(replica, partition) = |
| 884 | device_assignment[replica][partition]->id(); |
| 885 | } |
| 886 | } |
| 887 | return Compile(computation, std::move(argument_layouts), build_options, |
| 888 | std::move(client), xla_assignment); |
| 889 | } |
| 890 | |
| 891 | /*static*/ StatusOr<std::unique_ptr<PyLocalExecutable>> |
| 892 | PyLocalExecutable::Compile(const XlaComputation& computation, |
nothing calls this directly
no test coverage detected