<< replica_groups.getType().getRank();
  }
  return success();
}

ChannelHandleAttr ConvertChannel(OpBuilder& builder, int64_t channel_id,
                                 StringRef mode) {
  if (mode == "CrossReplica") {
    return ChannelHandleAttr();
  }
  return ChannelHandleAttr::get(builder.getContext(),
                                /*handle=*/channel_id,

                   StringRef key, size_t index, Value token,
                   StringRef host_handler_name, bool manual_sharding) {
  // type 2 == DEVICE_TO_HOST
  auto channel_handle = ChannelHandleAttr::get(builder.getContext(),
                                               /*handle=*/GetNextChannelId(),
                                               /*type=*/2);
  auto send = builder.create<SendOp>(

        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);