return false;
  }

  ModuleOp module = op->getParentOfType<ModuleOp>();
  func::FuncOp func = module.lookupSymbol<func::FuncOp>(
      op.getShapeInferenceGraphAttr().getRootReference());

  if (func == nullptr) return false;

  std::vector<_XlaRecvAtHostOp> xla_recv_at_host_ops;
  func.walk([&](_XlaRecvAtHostOp op) { xla_recv_at_host_ops.push_back(op); });