builder.setInsertionPoint(op);
  ChannelHandleAttr channel_handle = ConvertChannel(builder, channel_id, mode);
  Location loc = op->getLoc();
  Type element_type = getElementTypeOrSelf(input.getType());
  auto all_reduce = builder.create<AllReduceOp>(
      loc, result_type, input, replica_groups, channel_handle, nullptr);

  if (all_reduce.getNumResults() != 1) {
    return op->emitOpError()

  // type 2 == DEVICE_TO_HOST
  auto channel_handle = ChannelHandleAttr::get(builder.getContext(),
                                               /*handle=*/GetNextChannelId(),
                                               /*type=*/2);
  auto send = builder.create<SendOp>(
      loc, token.getType(), operand, token, channel_handle,
      /*is_host_transfer=*/builder.getBoolAttr(true));