| 263 | } |
| 264 | |
| 265 | StatusOr<std::unique_ptr<HloModuleConfig>> Service::CreateModuleConfig( |
| 266 | const ProgramShape& program_shape, |
| 267 | absl::Span<const Shape* const> argument_shapes, |
| 268 | const ExecutionOptions* execution_options, |
| 269 | const AotCompilationOptions* aot_options) { |
| 270 | auto config = absl::make_unique<HloModuleConfig>(program_shape); |
| 271 | ComputationLayout* computation_layout = |
| 272 | config->mutable_entry_computation_layout(); |
| 273 | if (program_shape.parameters_size() != argument_shapes.size()) { |
| 274 | return InvalidArgument("computation takes %d parameters, but %u given", |
| 275 | program_shape.parameters_size(), |
| 276 | argument_shapes.size()); |
| 277 | } |
| 278 | for (int i = 0; i < argument_shapes.size(); ++i) { |
| 279 | // Verify that shape of arguments matches the shape of the arguments in the |
| 280 | // ProgramShape. |
| 281 | if (!ShapeUtil::Compatible(*argument_shapes[i], |
| 282 | program_shape.parameters(i))) { |
| 283 | return InvalidArgument( |
| 284 | "Argument does not match shape of computation parameter %d: want " |
| 285 | "%s, got %s", |
| 286 | i, ShapeUtil::HumanString(program_shape.parameters(i)), |
| 287 | ShapeUtil::HumanString(*argument_shapes[i])); |
| 288 | } |
| 289 | TF_RETURN_IF_ERROR( |
| 290 | computation_layout->mutable_parameter_layout(i)->CopyLayoutFromShape( |
| 291 | *argument_shapes[i])); |
| 292 | } |
| 293 | if (execution_options != nullptr && |
| 294 | execution_options->has_shape_with_output_layout()) { |
| 295 | const Shape shape_with_output_layout( |
| 296 | execution_options->shape_with_output_layout()); |
| 297 | TF_RETURN_IF_ERROR( |
| 298 | ValidateResultShape(shape_with_output_layout, program_shape.result())); |
| 299 | TF_RETURN_IF_ERROR( |
| 300 | computation_layout->mutable_result_layout()->CopyLayoutFromShape( |
| 301 | shape_with_output_layout)); |
| 302 | } else { |
| 303 | // If the result layout is not set, then choose the default. |
| 304 | computation_layout->mutable_result_layout()->SetToDefaultLayout(); |
| 305 | } |
| 306 | |
| 307 | if (execution_options != nullptr) { |
| 308 | if (execution_options->num_replicas() > 0) { |
| 309 | config->set_replica_count(execution_options->num_replicas()); |
| 310 | } else { |
| 311 | config->set_replica_count(options_.number_of_replicas()); |
| 312 | } |
| 313 | if (execution_options->num_partitions() > 0) { |
| 314 | config->set_num_partitions(execution_options->num_partitions()); |
| 315 | } |
| 316 | config->set_seed(execution_options->seed()); |
| 317 | config->set_debug_options(execution_options->debug_options()); |
| 318 | } else { |
| 319 | config->set_replica_count(options_.number_of_replicas()); |
| 320 | config->set_debug_options(GetDebugOptionsFromFlags()); |
| 321 | } |
| 322 |
nothing calls this directly
no test coverage detected