| 283 | } |
| 284 | |
| 285 | StatusOr<HloInstruction*> MakeMapHlo(absl::Span<HloInstruction* const> operands, |
| 286 | HloComputation* map_computation) { |
| 287 | CHECK(!operands.empty()) << "Map Hlo requires at least one operand."; |
| 288 | HloComputation* computation = operands.front()->parent(); |
| 289 | std::vector<const Shape*> operand_shapes; |
| 290 | int64 max_operand_rank = 0; |
| 291 | for (const HloInstruction* operand : operands) { |
| 292 | CHECK_EQ(computation, operand->parent()); |
| 293 | operand_shapes.push_back(&operand->shape()); |
| 294 | max_operand_rank = std::max(max_operand_rank, operand->shape().rank()); |
| 295 | } |
| 296 | std::vector<int64> map_dims(max_operand_rank); |
| 297 | std::iota(map_dims.begin(), map_dims.end(), 0); |
| 298 | TF_ASSIGN_OR_RETURN( |
| 299 | Shape map_shape, |
| 300 | ShapeInference::InferMapShape( |
| 301 | operand_shapes, map_computation->ComputeProgramShape(), map_dims)); |
| 302 | return computation->AddInstruction( |
| 303 | HloInstruction::CreateMap(map_shape, operands, map_computation)); |
| 304 | } |
| 305 | |
| 306 | StatusOr<HloInstruction*> MakeReduceHlo(HloInstruction* operand, |
| 307 | HloInstruction* init_value, |