// needed.
    if (failed(SetCollectiveInfo(rewriter, replica_groups, all_reduce))) {
      return failure();
    }

    StringRef reduce_op = all_reduce.getReduceOp();

    StringRef merge_op, final_op;
    if (reduce_op == "Add") {
      merge_op = "Add";
      final_op = "Id";
    } else if (reduce_op == "Mul") {
      merge_op = "Mul";
      final_op = "Id";

        rewriter.getIntegerAttr(rewriter.getIntegerType(64),
                                scatter_dimension.getSExtValue()),
        replica_groups, ChannelHandleAttr());
    StringRef reduce_op = op.getReduceOp();
    if (reduce_op == "Add") {
      BuildReduceBody<AddOp>(element_type, &reduce_scatter.getComputation(),
                             &rewriter);
    } else if (reduce_op == "Mul") {