// CHECK:           %[[VAL_3:.*]] = "mhlo.scatter"(%[[VAL_0]], %[[VAL_1]], %[[VAL_2]]) <{
// CHECK-SAME:          indices_are_sorted = false, scatter_dimension_numbers = #mhlo.scatter<update_window_dims = [1], inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false
// CHECK-SAME:      }> ({
// CHECK:           ^bb0(%[[VAL_4:.*]]: tensor<f32>, %[[VAL_5:.*]]: tensor<f32>):

  // CHECK-SAME: indices_are_sorted = false,
  // CHECK-SAME: scatter_dimension_numbers =
  // CHECK-SAME:   update_window_dims = [2]
  // CHECK-SAME:   inserted_window_dims = [0]
  // CHECK-SAME:   scatter_dims_to_operand_dims = [0]
  // CHECK-SAME:   index_vector_dim = 2
  // CHECK-SAME: unique_indices = false
  // CHECK: ^{{.*}}([[LHS:%.*]]: tensor<f32>, [[RHS:%.*]]: tensor<f32>):

        op.getLoc(), output_type, init,
        GetI64ElementsAttr(output_shape, &rewriter));

    // Parameters for the generated scatter op.
    SmallVector<int64_t, 1> inserted_window_dims(1, 0);
    SmallVector<int64_t, 1> scatter_dims_to_operand_dims(1, 0);
    int64_t index_vector_dim = segment_ids_rank;

    // Put all parameters in a StructAttr.