target.addLegalDialect<arith::ArithDialect>();

  if (failed(applyPartialConversion(getOperation(), target,
                                    std::move(patterns)))) {
    getOperation().emitError("Composite lowering pass failed.");
    signalPassFailure();
  }
}

}  // namespace

// Creates an instance of the pass.
std::unique_ptr<OperationPass<ModuleOp>> CreateCompositeLoweringPass() {

        new_arg = block.addArgument(arg_type, loc);
        quantize_op.setOperand(new_arg);
      } else {
        input_type.print(llvm::errs() << "Requested input type ");
        quantize_op.emitError(" Couldn't be modified to the requested type.");
        return failure();
      }
      new_input_types[i] = arg_type;
      arg.dropAllUses();
      if (quantize_op.use_empty()) {
        quantize_op.erase();

  if (!tensorflow::ParseOutputArrayInfo(input_arrays_flag, &inputs).ok()) {
    return emitError(loc, "parsing input array info failed ")
               << input_arrays_flag,
           nullptr;
  }
  if (!tensorflow::ParseOutputArrayInfo(output_arrays_flag, &outputs).ok()) {
    return emitError(loc, "parsing output array info failed ")
               << output_arrays_flag,
           nullptr;
  }

  StringRef getDescription() const final {
    return "Renames the entrypoint in SavedModel to `main`";
  }

  void runOnOperation() override {
    auto fail = [&](Operation* op, std::string message) {
      op->emitError(message);
      signalPassFailure();
    };

    DenseMap<StringRef, func::FuncOp> entrypoints;
    auto module = getOperation();
    module.walk([&](func::FuncOp op) {

  if (keyword == "future") return FutureType::get(getContext());
  if (keyword == "promise") return PromiseType::get(getContext());
  if (keyword == "async_handle") return AsyncHandleType::get(getContext());

  parser.emitError(parser.getNameLoc(), "unknown type: ") << keyword;
  return mlir::Type();
}

// Print a type registered to this dialect.
void MlrtDialect::printType(mlir::Type type,

    mlir::DialectAsmParser &parser) const {
  llvm::StringRef keyword;
  if (parser.parseKeyword(&keyword)) return mlir::Type();

  if (keyword == "tensor") return TFTensorType::get(getContext());

  parser.emitError(parser.getNameLoc(), "unknown type: ") << keyword;
  return mlir::Type();
}

// Print a type registered to this dialect.
void TensorflowMlrtDialect::printType(mlir::Type type,

  ModuleOp module_op = getOperation();

  for (auto init_func_op : GetInitializerFunctions(module_op)) {
    if (!init_func_op) return;

    if (init_func_op.getBlocks().size() != 1) {
      init_func_op.emitError("expects exactly one block in the MLIR function");
      return signalPassFailure();
    }

    auto var_handle_ops =
        init_func_op.getBlocks().front().getOps<VarHandleOp>();

  for (auto func : getOperation().getOps<func::FuncOp>()) {
    auto walk_result = func.walk([&](func::CallOp op) -> WalkResult {
      if (!symbol_table.lookup<func::FuncOp>(op.getCallee()))
        return getOperation().emitError()
               << func.getName() << " is not in the funcs allowlist";
      return WalkResult::advance();
    });
    if (walk_result.wasInterrupted()) return signalPassFailure();
  }
}

}  // namespace

      module_op, /*initializer_type=*/kTfSavedModelInitializerRestoreType);
  if (init_func_op) {
    if (failed(
            applyPatternsAndFoldGreedily(init_func_op, std::move(patterns)))) {
      init_func_op->emitError(
          "Failed to remove variable assignment by const patterns.");
      signalPassFailure();
    }
  } else {
    LOG(INFO) << "Initializer function with type 'restore_op' does not exist. "

  auto module_op = getOperation();

  patterns.add<CastBf16OpsToF32>(ctx);
  populateWithGenerated(patterns);

  if (failed(applyPatternsAndFoldGreedily(module_op, std::move(patterns)))) {
    module_op.emitError() << "quant-cast-bf16-ops-to-f32 failed.";
    signalPassFailure();
  }
}

}  // namespace

// Creates an instance of the Cast BF16 ops to F32 pass.