| 305 | // reshape this result to [2,2,3] and finally transpose it to [2,3,2]. |
| 306 | |
| 307 | StatusOr<HloInstruction*> GatherExpander::ExpandInstruction( |
| 308 | HloInstruction* gather_instr) { |
| 309 | CHECK(!ShapeUtil::IsZeroElementArray(gather_instr->shape())); |
| 310 | |
| 311 | HloComputation* computation = gather_instr->parent(); |
| 312 | HloInstruction* operand = gather_instr->mutable_operand(0); |
| 313 | HloInstruction* start_indices = gather_instr->mutable_operand(1); |
| 314 | const Shape& start_indices_shape = start_indices->shape(); |
| 315 | const Shape& output_shape = gather_instr->shape(); |
| 316 | int64 output_rank = output_shape.dimensions_size(); |
| 317 | |
| 318 | const GatherDimensionNumbers& dim_numbers = |
| 319 | gather_instr->gather_dimension_numbers(); |
| 320 | |
| 321 | int64 gather_loop_trip_count = 1; |
| 322 | for (int64 i = 0, e = start_indices_shape.dimensions_size(); i < e; i++) { |
| 323 | if (i != dim_numbers.index_vector_dim()) { |
| 324 | gather_loop_trip_count *= start_indices_shape.dimensions(i); |
| 325 | } |
| 326 | } |
| 327 | |
| 328 | if (!IsInt32(gather_loop_trip_count)) { |
| 329 | return Unimplemented( |
| 330 | "Gather operations with more than 2147483647 gather indices are not " |
| 331 | "supported. This error occurred for %s.", |
| 332 | gather_instr->ToString()); |
| 333 | } |
| 334 | |
| 335 | TF_ASSIGN_OR_RETURN( |
| 336 | HloInstruction * canonical_start_indices, |
| 337 | CanonicalizeGatherIndices(start_indices, dim_numbers.index_vector_dim())); |
| 338 | |
| 339 | CHECK_EQ(gather_loop_trip_count, |
| 340 | canonical_start_indices->shape().dimensions(0)); |
| 341 | |
| 342 | HloInstruction* accumulator_init = CreateGatherLoopAccumulatorInitValue( |
| 343 | computation, output_shape.element_type(), |
| 344 | gather_instr->gather_slice_sizes(), gather_loop_trip_count, |
| 345 | gather_instr->gather_dimension_numbers()); |
| 346 | |
| 347 | StatusOr<std::vector<HloInstruction*>> gather_loop_result_or_error = |
| 348 | WhileUtil::MakeCountedLoop( |
| 349 | computation, gather_loop_trip_count, |
| 350 | {operand, canonical_start_indices, accumulator_init}, |
| 351 | [&](HloInstruction* indvar, |
| 352 | const std::vector<HloInstruction*>& loop_state) { |
| 353 | return GatherLoopBody(*gather_instr, indvar, loop_state); |
| 354 | }, |
| 355 | gather_instr->metadata()); |
| 356 | |
| 357 | TF_ASSIGN_OR_RETURN(std::vector<HloInstruction*> gather_loop_result, |
| 358 | gather_loop_result_or_error); |
| 359 | |
| 360 | HloInstruction* accumulator_result = gather_loop_result.back(); |
| 361 | |
| 362 | TF_ASSIGN_OR_RETURN( |
| 363 | HloInstruction* const accumulator_with_batch_dims_decanonicalized, |
| 364 | AdjustBatchDimsInAccumulator(start_indices->shape(), accumulator_result, |
nothing calls this directly
no test coverage detected