Expands or contracts the scatter indices in the updates tensor.
| 112 | |
| 113 | // Expands or contracts the scatter indices in the updates tensor. |
| 114 | static StatusOr<HloInstruction*> AdjustScatterDims( |
| 115 | const Shape& scatter_indices_shape, HloInstruction* updates, |
| 116 | int64 index_vector_dim) { |
| 117 | int64 num_scatter_dims = scatter_indices_shape.dimensions_size(); |
| 118 | if (index_vector_dim < scatter_indices_shape.dimensions_size()) { |
| 119 | --num_scatter_dims; |
| 120 | } |
| 121 | if (num_scatter_dims == 0) { |
| 122 | // If there are no scatter dims, this must be a dynamic-update-slice kind of |
| 123 | // scatter. In this case, we prepend a degenerate dimension to work |
| 124 | // uniformly in the while loop. |
| 125 | return PrependDegenerateDims(updates, 1); |
| 126 | } |
| 127 | return CollapseFirstNDims(updates, num_scatter_dims); |
| 128 | } |
| 129 | |
| 130 | // Expands an index vector from the scatter_indices tensor into a vector that |
| 131 | // can be used to dynamic-update-slice to perform the scatter update. |
no test coverage detected