it->second.data_type;
      }
    }
  }
  func_op.eraseArguments(indices_to_erase);
  func_op.setType(
      FunctionType::get(func_op.getContext(), new_types,
                        llvm::to_vector<4>(return_op->getOperandTypes())));
}

// Lifts reads/writes of resource arguments from func_op and changes its
// signature. resource_data_types is the (index, data type) pair for each

  auto out_types = GetValueTypes(outputs);
  builder.setInsertionPointToEnd(&module.getBodyRegion().back());
  auto func_op = builder.create<func::FuncOp>(
      module.getLoc(), name, builder.getFunctionType(in_types, out_types));
  func_op.setPrivate();
  symbol_table.insert(func_op);
  return func_op;
}

TF::StatefulPartitionedCallOp EncapsulateOpsInFunc(
    OpBuilder& builder, SymbolTable& symbol_table,

  // Use, run, and verify
  TF_Operation* two = ScalarConst(2, host_graph_, s_);
  TF_Operation* func_feed = Placeholder(host_graph_, s_);
  TF_Operation* func_op = Use({two, func_feed});
  Run({{func_feed, Int32Tensor(3)}}, {{func_op, 0}, {func_op, 1}}, {3, 2});
  VerifyFDef(empty_, M({{"feed1_0"}, {"feed2_0"}}), M({{"feed2"}, {"feed1"}}),
             {{"feed1_0", "feed1"}, {"feed2_0", "feed2"}}, {});
}

        /*is_hybrid=*/is_hybrid);

    const mlir::func::FuncOp quantized_func = dyn_cast_or_null<func::FuncOp>(
        symbol_table.lookup(quantized_function_name));
    if (quantized_func == nullptr) {
      call_op->emitError("Failed to find the quantized function: " +
                         quantized_function_name);
      return failure();
    }
    mlir::func::FuncOp new_quantized_func =

        queue.push_back(op);
      }
    }
  }

  return visited;
}

// We want to adjust the func op according to some cross ops information.
static StatusOr<FuncOp> PostProcessFuncOp(FuncOp func) {
  OpBuilder builder(func);
  // When a quantized constant is imported, its quantization parameter is set
  // to be narrow range. Here revert to be the fully range if the user doesn't

    then_branch_op.setVisibility(func::FuncOp::Visibility::Private);

    // Constructs `else_branch`, which is executed when `if_cond` evaluates to
    // false.
    auto else_branch_op =
        rewriter.create<func::FuncOp>(loc, "cond_false", func_type);
    CreateCondFalseBranch(loc, shape_dtype, result_type, else_branch_op,
                          &rewriter);
    else_branch_op.setVisibility(func::FuncOp::Visibility::Private);

    func::FuncOp ResolveThenFunction(::mlir::SymbolTableCollection* table) {
      if (table)
        return table->lookupNearestSymbolFrom<func::FuncOp>(*this, getThenBranchAttr());
      return SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(
        *this, getThenBranchAttr());
    }
    // TODO(b/204997177): Deprecate and remove.

  output_types.reserve(outputs.size());
  for (Value v : outputs) output_types.emplace_back(v.getType());

  auto func_type = builder->getFunctionType(operand_types, output_types);

  FuncOp outlined_func =
      FuncOp::create(ops.front()->getLoc(), kHostFunctionAttr, func_type);

  // Create function body.
  Block* outlined_func_block = outlined_func.addEntryBlock();

// TF2XLA ops aren't supported by later stages.
LogicalResult ConvertTf2XlaOps(func::FuncOp func, MLIRContext *context) {
  ConversionTarget target(*context);
  target.addLegalDialect<arith::ArithDialect>();
  target.addLegalDialect<func::FuncDialect>();
  target.addLegalDialect<TF::TensorFlowDialect>();
  target.addLegalOp<ModuleOp>();
  target.addLegalOp<func::FuncOp>();
  target.addIllegalOp<TF::XlaConvV2Op>();

#include <cstdint>
#include <memory>
#include <utility>

#include "absl/algorithm/container.h"
#include "llvm/ADT/Sequence.h"
#include "llvm/Support/Debug.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/Dialect/Quant/QuantOps.h"  // from @llvm-project  // NOLINT: Required to register quantization dialect.
#include "mlir/Dialect/Quant/QuantTypes.h"  // from @llvm-project