// Result types.
    SmallVector<Type, 4> result_types;
    for (auto result_type : op->getResultTypes()) {
      if (IsQI8Type(result_type) || IsQUI8Type(result_type)) {
        auto dequantized_result_type =
            mlir::quant::QuantizedType::castToExpressedType(result_type);
        result_types.push_back(dequantized_result_type);
      } else {
        result_types.push_back(result_type);
      }
    }

  CHECK(func);
  CHECK_EQ(func->result_types().size(), 113);
  CHECK_EQ(func->argument_types().size(), 113);

  auto arguments = CreateTestArguments(func, host);

  tfrt::ResourceContext resource_ctx;
  auto exec_ctx = tfd::CreateFallbackTestExecutionContext(host, &resource_ctx);

  std::vector<tfrt::RCReference<tfrt::AsyncValue>> results;
  results.resize(func->result_types().size());

      mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
  mlir::OpBuilder builder(module_ref->getBodyRegion());

  llvm::SmallVector<mlir::Type, 8> result_types;
  auto cluster = builder.create<mlir::tf_device::ClusterOp>(
      mlir::UnknownLoc::get(&context), result_types);
  cluster->setAttr(kNumCoresPerReplicaAttr,
                   builder.getIntegerAttr(builder.getIntegerType(64), 1));

  // Create attributes used for creating an XlaCallModuleOp.
  SmallVector<Attribute> shape_attrs;
  for (const Type result_type : result_types) {
    shape_attrs.push_back(
        tf_type::ShapeAttr::get(ctx, mlir::cast<ShapedType>(result_type)));
  }
  const auto empty_array_attr = ArrayAttr::get(ctx, {});
  // TODO: b/310291615 - find a better way for platform support.

// apart from the arg and retval nodes corresponding to `args` and
// `result_types` respectively.
absl::StatusOr<std::unique_ptr<Graph>> CreateSingleOpGraph(
    const NodeDef& node_def, absl::Span<const XlaArgument> args,
    absl::Span<const DataType> result_types);

// Checks if single device compilation and execution with PJRT is enabled for

  Operation* last_cluster_op = cluster_ops.back();
  OpBuilder builder(last_cluster_op);

  llvm::SmallVector<Type, 8> result_types;
  for (Value result : results) result_types.push_back(result.getType());
  auto cluster = builder.create<mlir::tf_device::ClusterOp>(
      last_cluster_op->getLoc(), result_types);

  Block* body = new Block;
  cluster.getBody().push_back(body);

    SmallVector<Type, 4> result_types;
    for (Value result : call_op->getResults()) {
      TensorType result_type = mlir::dyn_cast<TensorType>(result.getType());
      if (!result_type) {
        result_types.push_back(result.getType());
        continue;
      }
      QuantizedType qtype =
          mlir::dyn_cast<QuantizedType>(result_type.getElementType());
      if (!qtype) {

    if (IsOnLocalHost(host)) continue;

    llvm::SmallVector<mlir::Type, 4> input_types;
    llvm::SmallVector<mlir::Type, 4> result_types;
    for (Value input : metadata.inputs) {
      input_types.push_back(input.getType());
    }
    for (Value result : metadata.results) {
      result_types.push_back(result.getType());
    }

    // Replaces ':' and '/' with '_' in the host name and uses the resulting

    std::vector<mlir::Type> result_types;

    for (mlir::Type t : op.getResultTypes()) {
      auto ty = mlir::cast<mlir::TensorType>(t);
      if (!ty.hasStaticShape()) return mlir::WalkResult::interrupt();
      result_types.push_back(t);
    }

    auto layout = mlir::GetTPUInfeedLayout(
        mlir::TupleType::get(builder.getContext(), result_types), builder);

  for (int64_t i = 0, end = result_types.size(); i < end; ++i) {
    Node* node;
    string retval_name = absl::StrCat("_retval", i);
    Status status = NodeBuilder(retval_name, FunctionLibraryDefinition::kRetOp)
                        .Input(main_node, i)
                        .Attr("T", result_types[i])
                        .Attr("index", i)