message: Improve code generated by scatter expander
  After this change if the indexes dimension is 1 then we remove it and
  treat the scatter as a scatter with scalar indices.
diff: @@ -26,7 +26,6 @@ limitations under the License.
  namespace xla {
  -
  // Transposes the given scatter_indices such that the index_vector_dim becomes
  // the most-minor dimension.
  static StatusOr<HloInstruction*> TransposeIndexVectorDimToLast(
  @@ -60,6 +59,13 @@ static StatusOr<HloInstruction*> CanonicalizeScatterIndices(
  TF_ASSIGN_OR_RETURN(
  HloInstruction * transposed_scatter_indices,
  TransposeIndexVectorDimToLast(scatter_indices, index_vector_dim));
  + if (ShapeUtil::Rank(scatter_indices->shape()) == index_vector_dim + 1 &&
  + scatter_indices->shape().dimensions(index_vector_dim) == 1) {
  + auto new_shape =
  + ShapeUtil::DeleteDimension(index_vector_dim, scatter_indices->shape());
  + TF_ASSIGN_OR_RETURN(scatter_indices,
  + MakeReshapeHlo(new_shape, scatter_indices));
  + }
  bool indices_are_scalar =
  index_vector_dim == scatter_indices->shape().dimensions_size();
  @@ -214,9 +220,6 @@ static StatusOr<std::vector<HloInstruction*>> ScatterLoopBody(
  HloInstruction* updates = loop_state[2];
  bool has_scalar_indices = scatter_indices->shape().dimensions_size() == 1;
  - CHECK_EQ(has_scalar_indices,
  - dim_numbers.index_vector_dim() ==
  - scatter->operand(1)->shape().dimensions_size());
  // Build a vector form of the induction variable of the while loop.
  TF_ASSIGN_OR_RETURN(
  