auto new_key = StringAttr::get(ctx, old_key.str() + "_dtoh_0");
      op->setAttr("key", new_key);
    } else if (isa<TF::XlaRecvFromHostOp>(op)) {
      StringRef old_key = op->getAttrOfType<StringAttr>("key").getValue();
      auto new_key = StringAttr::get(ctx, old_key.str() + "_htod_0");
      op->setAttr("key", new_key);
    }
  });
  return success();
}

  MLIRContext* context = op->getContext();

  std::string formatted_key =
      device_to_host ? llvm::formatv("{0}_dtoh_{1}", key, index).str()
                     : llvm::formatv("{0}_htod_{1}", key, index).str();

  auto rendezvous_name = StringAttr::get(context, formatted_key);
  auto rendezvous_name_attr = NamedAttribute(
      StringAttr::get(context, xla::kXlaHostTransferRendezvousNameAttr),

    Graph* g) {
  NodeDefBuilder send_pred_builder(name, "XlaSendToHost");
  send_pred_builder.Attr("Tinput", DT_BOOL);
  send_pred_builder.Attr("key", absl::StrCat(host_transfer_key, "_dtoh_0"));
  send_pred_builder.Attr(kXlaTokenInputNodesAttrName,
                         std::vector<string>{kXlaTokenArgNodeName});
  send_pred_builder.Attr(kXlaOriginalOutsideCompilationNodeName, name);