func.walk([&](Operation* op) {
    b.setInsertionPoint(op);
    if (auto dq = llvm::dyn_cast<DequantizeOp>(op)) {
      auto dcast = b.create<quantfork::DequantizeCastOp>(
          dq.getLoc(), dq.getOutput().getType(), dq.getInput());
      dq.getOutput().replaceAllUsesWith(dcast);
      dq.erase();
    } else if (auto q = llvm::dyn_cast<QuantizeOp>(op)) {
      auto qcast = b.create<quantfork::QuantizeCastOp>(

  const SavedObjectGraph object_graph = tf_package->GetObjectGraph();
  auto& nodes = object_graph.nodes();
  // Get a map of the concrete functions to their input / output signatures.
  auto& concrete_functions = object_graph.concrete_functions();
  auto& root = nodes.at(0);
  for (auto& child : root.children()) {
    // The child's name describes the name of the edge that connects to the

  int num_threads = 3;

  TensorShape data_shape({1, 1});
  auto a = Placeholder(root, DT_INT32, Placeholder::Shape(data_shape));
  auto b = Placeholder(root, DT_INT32, Placeholder::Shape(data_shape));
  auto c = BatchMatMul(root, a, b);
  ClientSession session(root);
  std::vector<Tensor> outputs;

  auto inter_op_threadpool =
      absl::make_unique<CustomThreadPoolImpl>(num_threads);

}

bool checkWhetherGraphHasValidStaticLookupTables(ModuleOp module) {
  auto hashtables = GetAllOps<TF::HashTableV2Op>(&module);
  // No needs to run the legalization patterns.
  if (hashtables.empty()) {
    return false;
  }

  for (auto hashtable : hashtables) {
    auto key_dtype = hashtable.getKeyDtype();
    auto value_dtype = hashtable.getValueDtype();

  auto in_spec = input_specs[0].dyn_cast<UniformQuantizedType>();
  // TODO(fengliuai): handles the PerAxis QuantizedType.
  auto w_spec = input_specs[1].dyn_cast<UniformQuantizedType>();
  auto b_spec = input_specs[2].dyn_cast<UniformQuantizedType>();
  auto o_spec = out_specs[0].dyn_cast<UniformQuantizedType>();
  if (!in_spec || !w_spec || !b_spec || !o_spec) return failure();

  double scale_product = in_spec.getScale() * w_spec.getScale();

void PopulateTfVersions(mlir::ModuleOp module, const VersionDef& versions) {
  mlir::Builder b(module.getContext());
  auto producer =
      b.getNamedAttr("producer", b.getI32IntegerAttr(versions.producer()));
  auto min_consumer = b.getNamedAttr(
      "min_consumer", b.getI32IntegerAttr(versions.min_consumer()));
  auto bad_consumers = b.getNamedAttr(
      "bad_consumers",
      b.getI32ArrayAttr(llvm::ArrayRef<int32_t>(

                                 ArrayRef<StateT *> results) {
    if (auto cast = dyn_cast<TF::CastOp>(op)) {
      this->join(results[0], *operands[0]);
    } else if (auto while_op = dyn_cast<TF::WhileRegionOp>(op)) {
      for (auto &region : while_op->getRegions()) {
        for (auto [arg, value] :
             llvm::zip(region.getArguments(), while_op->getOperands())) {

}

TEST(BatchFunctionTest, Basic) {
  auto corert = CreateTestCoreRuntime();
  auto* host = corert->GetHostContext();
  auto [bef_buffer, bef_file] = CreateBefFile(host);
  auto* func = bef_file->GetFunction("main");
  CHECK(func);
  CHECK_EQ(func->result_types().size(), 113);
  CHECK_EQ(func->argument_types().size(), 113);

  auto arguments = CreateTestArguments(func, host);

func::FuncOp createMaxUnpoolingFunc(
    mlir::Builder* builder, const SmallVector<int64_t, 4>& input_shape,
    const SmallVector<int64_t, 4>& output_shape) {
  auto input_type = RankedTensorType::get(input_shape, builder->getF32Type());
  auto indices_type = RankedTensorType::get(input_shape, builder->getI64Type());
  auto output_type = RankedTensorType::get(output_shape, builder->getF32Type());
  SmallVector<mlir::Type, 2> input_types{input_type, indices_type};

  const auto num_replicas = replicate_op.getN();
  auto return_op = llvm::dyn_cast<tf_device::ReturnOp>(
      replicate_op.getRegion().front().getTerminator());

  // Get the new result types.
  // TODO(prakalps): Do not add a value to returned values if it is already
  // returned.
  auto new_result_types = llvm::to_vector<4>(replicate_op->getResultTypes());
  for (auto assign : tail_assign_variable_ops) {