}

  pm.addNestedPass<mlir::func::FuncOp>(mlir::TF::CreateLowerQuantizedPass());
  pm.addNestedPass<mlir::func::FuncOp>(
      mlir::quant::stablehlo::CreateConvertTFQuantTypesPass());

  if (lower_to_xla_hlo) {
    for (auto& target_pass : custom_legalization_passes) {
      pm.addNestedPass<mlir::func::FuncOp>(std::move(target_pass));
    }

  // ranges.
  bool SetInputNodesQuantizationParams(func::FuncOp func);

  // The function might contain more stats ops than required, and it will
  // introduce requantize if the calibration stats have conflicts. This method
  // tries to remove all the redundant stats ops.
  bool RemoveRedundantStats(func::FuncOp func);

  // Verify the quantization specification is expected for quantizing the

  CopyDeviceAndUnderscoredAttributes(generator_op, generator_region);

  func::FuncOp init_function =
      SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(
          generator_op, generator_op.getInitFunc());
  func::FuncOp next_function =
      SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(
          generator_op, generator_op.getNextFunc());
  func::FuncOp finalize_function =
      SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(

    }

    return absl::OkStatus();
  }

  mlir::func::FuncOp GetMainFunc() {
    func::FuncOp main_func = module_->lookupSymbol<mlir::func::FuncOp>("main");
    EXPECT_TRUE(main_func);

    return main_func;
  }

  mlir::Operation& GetFirstOpFromMain() {
    mlir::func::FuncOp main_func = GetMainFunc();
    return main_func.getBody().front().front();
  }

    // and use it for `shape_inference_graph` attribute on XlaHostCompute.
    func::FuncOp cloned_func;
    SymbolTable manager(op->getParentOfType<ModuleOp>());
    StringRef host_module = op.getHostMlirModule();
    if (!host_module.empty()) {
      mlir::OwningOpRef<mlir::ModuleOp> module_for_func;

      func::FuncOp func = op.GetHostFunc(&module_for_func);

      OpBuilder::InsertionGuard guard(rewriter);

  for (const Value output_value : outputs) {
    result_types.push_back(output_value.getType());
  }

  // 1) Create FuncOp for the StableHLO ops. They will be separate subgraphs.
  builder.setInsertionPoint(&*module_op.begin());
  auto stablehlo_func_op = builder.create<func::FuncOp>(
      module_op.getLoc(), CreateStablehloFunctionName(stablehlo_func_id),
      FunctionType::get(ctx, arg_types, result_types));

  // ranges.
  bool SetInputNodesQuantizationParams(func::FuncOp func);

  // The function might contain more stats ops than required, and it will
  // introduce requantize if the calibration stats have conflicts. This method
  // tries to remove all the redundant stats ops.
  bool RemoveRedundantStats(func::FuncOp func);

  // Verify the quantization specification is expected for quantizing the

    ModuleOp module) {
  SymbolTableCollection table;
  SymbolUserMap symbol_map(table, module);
  llvm::DenseMap<func::FuncOp, llvm::DenseSet<func::FuncOp>> caller_callee_map;
  // Creates work queue for determining reachability below.
  std::queue<func::FuncOp> function_worklist;

  for (auto func : module.getOps<func::FuncOp>()) {
    for (auto user : symbol_map.getUsers(func)) {

}

// Returns the entry function, i.e. the callee of `xla_call_module_op`.
func::FuncOp GetEntryFuncOp(TF::XlaCallModuleOp xla_call_module_op,
                            const SymbolTable symbol_table) {
  const auto entry_function_symbol_ref =
      xla_call_module_op->getAttrOfType<FlatSymbolRefAttr>(kEntryFuncAttrName);

  return dyn_cast_or_null<func::FuncOp>(
      symbol_table.lookup(entry_function_symbol_ref.getValue()));
}

  void TearDown() override {
    fused_lstm_func_.erase();
    fused_lstm_func_cifg_.erase();
    fused_ln_lstm_func_.erase();
    builder_.reset();
  }

  func::FuncOp fused_lstm_func_;
  func::FuncOp fused_lstm_func_cifg_;
  func::FuncOp fused_ln_lstm_func_;
  std::unique_ptr<mlir::MLIRContext> context_;
  std::unique_ptr<mlir::Builder> builder_;
};

TEST_F(LstmUtilsTest, ConvertLSTMCellSimple) {