let dependentDialects = ["mhlo::MhloDialect", "sparse_tensor::SparseTensorDialect"];
}

def VerifyTFXLALegalization : Pass<"tfxla-verify-legalization", "mlir::func::FuncOp"> {
  let summary = "Verifies that all TF ops have been legalized to XLA.";

  let description = [{"Ensures that all Tensorflow ops have been legalized to "

namespace mlir {
namespace TF {

// Legalize the tf ops to the quant ops, so the quantization passes can work.
std::unique_ptr<OperationPass<func::FuncOp>> CreateLegalizeTFToQuantPass();

// Fallbacks ops that are not supported by TF Quantization to TFLite Flex ops.
std::unique_ptr<OperationPass<func::FuncOp>> CreateFallbackToFlexOpsPass(
    const std::string &mode = "DEFAULT");

}  // namespace TF
}  // namespace mlir

    // Note: lookupSymbol is a linear lookup which means the overall
    // complexity = # ReduceDataset ops x # of functions in module.
    func::FuncOp reduce_func =
        function->getParentOfType<ModuleOp>().lookupSymbol<FuncOp>(
            reduce_dataset.getF());

    // The reduce function arguments consist of three part in this order:
    // 1. Reduction state inputs.
    // 2. Dataset inputs.

  // Move all remaining nested functions back into the parent module.
  Block &nested_block = nested_module->getRegion(0).front();
  for (func::FuncOp func_op :
       llvm::make_early_inc_range(nested_block.getOps<func::FuncOp>())) {
    if (!symbol_table.lookupSymbolIn(getOperation(), func_op.getName())) {
      nested_block.getOperations().remove(func_op.getOperation());

    FunctionType func_type =
        FunctionType::get(context, input_types, result_types);
    Location loc = metadata.ops.front()->getLoc();
    func::FuncOp func_op = func::FuncOp::create(loc, func_name, func_type);
    // Sets the device attribute for every input and every result of the
    // function.
    for (int i : llvm::seq<int>(0, metadata.input_devices.size())) {

    // side effect analysis of their parent functions are invalidated. They
    // need to be recomputed.
    recompute_analysis_for_funcs.insert(
        while_op->getParentOfType<func::FuncOp>());
    func::FuncOp while_cond = while_op.cond_function();
    // Rewrite while op with extra chaining arguments and results.
    while_op = RewriteWhileOp(while_op, num_chains, chaining_data_type);

  if (!status.ok()) {
    return op.emitError()
           << "attribute 'host_mlir_module' can not be deserialized. "
           << status.message();
  }

  func::FuncOp func = module_for_func->lookupSymbol<func::FuncOp>("host_func");
  if (!func)
    return op.emitError()
           << "serialized module in attribute 'host_mlir_module' does not "
              "contain 'host_func' function.";

    mlir::func::FuncOp func_op) {
  for (mlir::Operation& op : func_op.front()) {
    if (!mlir::isMemoryEffectFree(&op) && !IsTensorArrayOp(&op)) return false;
  }

  return true;
}

TensorArraySideEffectAnalysis::TensorArraySideEffectAnalysis(
    mlir::ModuleOp module) {
  for (auto func_op : module.getOps<mlir::func::FuncOp>()) {

        parseSourceString<ModuleOp>(module_string, op.getContext());
    if (!nested_module) {
      op->emitError() << "could not parse attached MLIR module";
      return;
    }
    func::FuncOp main =
        dyn_cast<func::FuncOp>(nested_module->lookupSymbol("main"));
    if (!main) {
      op->emitError() << "MLIR Opaque Op expects a main() entry point\n";
      return;
    }

                           ->getAttrOfType<FlatSymbolRefAttr>(
                               TF::kStablehloEntryFunctionAttrName)
                           .getValue();
  func::FuncOp func_op = symbol_table.lookup<func::FuncOp>(function_name);

  // We should not unwrap if the function is not from
  // ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass.
  if (!func_op->hasAttr(TF::kFromXlaCallModuleAttrName)) {