using llvm::DefInit;
using llvm::dyn_cast;
using llvm::formatv;
using llvm::LessRecord;
using llvm::raw_ostream;
using llvm::Record;
using llvm::RecordKeeper;
using llvm::RecordRecTy;
using llvm::SmallVector;
using llvm::StringInit;
using llvm::StringRef;

enum ActionType {
  OpConv,
  RuntimeVerify,
};

// NOLINTNEXTLINE
llvm::cl::opt<ActionType> action(
    llvm::cl::desc("Action to perform:"),

#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/IR/Builders.h"  // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project

    func::FuncOp float_func =
        dyn_cast<func::FuncOp>(symbol_table.lookup(function_name));
    OperandRange func_args = op.getArgs();
    func::FuncOp new_float_func = float_func.clone();

    SmallVector<Value> new_float_func_args{func_args.begin(), func_args.end()};
    new_float_func_args[weight_operand_idx] = reshape_op;
    new_float_func.getArgument(weight_operand_idx).setType(new_shape);

  BlockArgument new_file_prefix_arg =
      main_func_op.getBody().front().addArgument(
          file_prefix_arg_type,
          NameLoc::get(builder.getStringAttr(kTfFilePrefix)));

  SmallVector<Type> input_types(main_func_op.getArgumentTypes());
  input_types.emplace_back(file_prefix_arg_type);

  main_func_op.setType(
      builder.getFunctionType(input_types, main_func_op.getResultTypes()));

      return failure();
    }
    return success();
  }
  void rewrite(TF::_XlaHostComputeMlirOp op,
               PatternRewriter& rewriter) const override {
    llvm::SmallVector<Attribute> shape_attrs;
    shape_attrs.reserve(op.getNumResults());
    for (Type ty : op.getResultTypes()) {
      shape_attrs.push_back(TF::ShapeAttr::get(rewriter.getContext(),

#include <optional>
#include <string>
#include <utility>
#include <vector>

#include "absl/container/flat_hash_map.h"
#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
#include "llvm/ADT/SmallVector.h"
#include "llvm/Analysis/AssumeBundleQueries.h"
#include "mlir/IR/Attributes.h"  // from @llvm-project
#include "mlir/IR/Builders.h"  // from @llvm-project
#include "mlir/IR/Operation.h"  // from @llvm-project

#ifndef TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_QUANTIZATION_CONTEXT_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_QUANTIZATION_CONTEXT_H_

#include <vector>

#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallVector.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/Dialect/Quant/QuantTypes.h"  // from @llvm-project
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project

#include <optional>
#include <string>
#include <vector>

#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project

#include "absl/strings/str_cat.h"
#include "llvm/ADT/APFloat.h"
#include "llvm/ADT/APInt.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/Sequence.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/StringSwitch.h"
#include "llvm/ADT/iterator_range.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/FormatVariadic.h"

    bool WrapsSingleOp();
  }];

  let builders = [
    OpBuilder<(ins "int":$n,
      "const llvm::SmallDenseMap<StringRef, llvm::SmallVector<StringRef, 4>>&":$devices,
      "llvm::ArrayRef<std::pair<ValueRange, Type>>":$replicated_inputs,
      "ValueRange":$packed_inputs, "TypeRange":$replica_output_types)>,