| 343 | // f. Write the updated value of the slice into the operand tensor. |
| 344 | |
| 345 | StatusOr<HloInstruction*> ScatterExpander::ExpandScatter( |
| 346 | HloInstruction* scatter) { |
| 347 | HloInstruction* operand = scatter->mutable_operand(0); |
| 348 | HloInstruction* scatter_indices = scatter->mutable_operand(1); |
| 349 | HloInstruction* updates = scatter->mutable_operand(2); |
| 350 | const ScatterDimensionNumbers& dim_numbers = |
| 351 | scatter->scatter_dimension_numbers(); |
| 352 | |
| 353 | // If the updates tensor is empty, there is no need to update the operand. We |
| 354 | // can return the operand as is. |
| 355 | if (ShapeUtil::IsZeroElementArray(updates->shape())) { |
| 356 | return operand; |
| 357 | } |
| 358 | |
| 359 | // Compute the trip count for the while loop to be used for scatter. This |
| 360 | // should be the number of indices we should scatter into the operand. |
| 361 | const Shape& scatter_indices_shape = scatter_indices->shape(); |
| 362 | int64 scatter_loop_trip_count = 1; |
| 363 | for (int64 i = 0, e = scatter_indices_shape.dimensions_size(); i < e; i++) { |
| 364 | if (i != dim_numbers.index_vector_dim()) { |
| 365 | scatter_loop_trip_count *= scatter_indices_shape.dimensions(i); |
| 366 | } |
| 367 | } |
| 368 | if (!IsInt32(scatter_loop_trip_count)) { |
| 369 | return Unimplemented( |
| 370 | "Scatter operations with more than 2147483647 scatter indices are not " |
| 371 | "supported. This error occurred for %s.", |
| 372 | scatter->ToString()); |
| 373 | } |
| 374 | |
| 375 | // Canonicalize the scatter_indices, after which the size of its most-major |
| 376 | // dimension must be same as the while loop trip count. |
| 377 | TF_ASSIGN_OR_RETURN(HloInstruction * canonical_scatter_indices, |
| 378 | CanonicalizeScatterIndices( |
| 379 | scatter_indices, dim_numbers.index_vector_dim())); |
| 380 | CHECK_EQ(scatter_loop_trip_count, |
| 381 | canonical_scatter_indices->shape().dimensions(0)); |
| 382 | |
| 383 | // Canonicalize the updates, after which the size of its most-major dimension |
| 384 | // must be same as the while loop trip count. |
| 385 | TF_ASSIGN_OR_RETURN( |
| 386 | HloInstruction * canonical_updates, |
| 387 | PermuteScatterAndWindowDims( |
| 388 | updates, AsInt64Slice(dim_numbers.update_window_dims()))); |
| 389 | TF_ASSIGN_OR_RETURN( |
| 390 | HloInstruction * adjusted_canonical_updates, |
| 391 | AdjustScatterDims(scatter_indices->shape(), canonical_updates, |
| 392 | dim_numbers.index_vector_dim())); |
| 393 | CHECK_EQ(scatter_loop_trip_count, |
| 394 | adjusted_canonical_updates->shape().dimensions(0)); |
| 395 | |
| 396 | // The while loop that implements the scatter operation. |
| 397 | StatusOr<std::vector<HloInstruction*>> scatter_loop_result_status = |
| 398 | WhileUtil::MakeCountedLoop( |
| 399 | scatter->parent(), scatter_loop_trip_count, |
| 400 | {operand, canonical_scatter_indices, adjusted_canonical_updates}, |
| 401 | [&](HloInstruction* induction_var, |
| 402 | const std::vector<HloInstruction*>& loop_state) { |
nothing calls this directly
no test coverage detected