loc, op.getType(), op.getInput(),
        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);