static*/
| 889 | } |
| 890 | |
| 891 | /*static*/ StatusOr<std::unique_ptr<PyLocalExecutable>> |
| 892 | PyLocalExecutable::Compile(const XlaComputation& computation, |
| 893 | absl::optional<std::vector<Shape>> argument_layouts, |
| 894 | const ExecutableBuildOptions* build_options, |
| 895 | std::shared_ptr<PyLocalClient> client, |
| 896 | absl::optional<DeviceAssignment> device_assignment) { |
| 897 | tensorflow::profiler::TraceMe traceme("LocalExecutable::Compile"); |
| 898 | |
| 899 | ExecutableBuildOptions options; |
| 900 | if (build_options != nullptr) { |
| 901 | options = *build_options; |
| 902 | } |
| 903 | |
| 904 | if (!options.device_allocator()) { |
| 905 | options.set_device_allocator(client->allocator()); |
| 906 | } |
| 907 | |
| 908 | if (device_assignment) { |
| 909 | VLOG(2) << "PyLocalExecutable::Compile got device_assignment:\n" |
| 910 | << device_assignment->ToString(); |
| 911 | if (device_assignment->replica_count() != options.num_replicas()) { |
| 912 | return InvalidArgument( |
| 913 | "Mismatched number of replicas for device " |
| 914 | "assignment and computation (%d vs %d).\n%s", |
| 915 | device_assignment->replica_count(), options.num_replicas(), |
| 916 | device_assignment->ToString()); |
| 917 | } |
| 918 | if (device_assignment->computation_count() != options.num_partitions()) { |
| 919 | return InvalidArgument( |
| 920 | "Mismatched number of partitions for device " |
| 921 | "assignment and computation (%d vs %d).\n%s", |
| 922 | device_assignment->computation_count(), options.num_partitions(), |
| 923 | device_assignment->ToString()); |
| 924 | } |
| 925 | } else { |
| 926 | TF_ASSIGN_OR_RETURN(device_assignment, |
| 927 | client->GetDefaultDeviceAssignment( |
| 928 | options.num_replicas(), options.num_partitions())); |
| 929 | VLOG(2) << "PyLocalExecutable::Compile using default device_assignment:\n" |
| 930 | << device_assignment->ToString(); |
| 931 | } |
| 932 | options.set_device_assignment(device_assignment.value()); |
| 933 | |
| 934 | if (!argument_layouts) { |
| 935 | TF_ASSIGN_OR_RETURN(ProgramShape program_shape, |
| 936 | computation.GetProgramShape()); |
| 937 | argument_layouts = program_shape.parameters(); |
| 938 | for (Shape& shape : *argument_layouts) { |
| 939 | LayoutUtil::ClearLayout(&shape); |
| 940 | } |
| 941 | } |
| 942 | std::vector<const Shape*> argument_layout_pointers; |
| 943 | argument_layout_pointers.reserve(argument_layouts->size()); |
| 944 | |
| 945 | // Assign a default layout to any array subshapes that are missing layouts. |
| 946 | auto assign_layouts = [client](Shape* shape) { |
| 947 | return ShapeUtil::ForEachMutableSubshapeWithStatus( |
| 948 | shape, [&](Shape* subshape, const ShapeIndex&) { |