public:
  // Returns std::nullopt if the operation cannot be folded.
  static std::optional<FoldAdaptor> Create(Operation* operation) {
    auto foldable_opr = [](Value val) -> bool {
      return !llvm::isa<BlockArgument>(val) &&
             llvm::isa<stablehlo::ConstantOp>(val.getDefiningOp());
    };
    if (!llvm::all_of(operation->getOperands(), foldable_opr)) {
      return std::nullopt;
    }
    return FoldAdaptor(operation);

using OpShardingVector = llvm::SmallVector<OpShardingVariant, 8>;
using OptionalOpShardingVector =
    llvm::SmallVector<std::optional<OpShardingVariant>, 8>;
using llvm::StringRef;
using mlir::Block;
using mlir::BlockArgument;
using mlir::BoolAttr;
using mlir::Builder;
using mlir::IntegerAttr;
using mlir::LogicalResult;
using mlir::ModuleOp;
using mlir::Operation;
using mlir::OpOperand;
using mlir::OpResult;

          continue;
        }
        // Transfer the annotation to both region arguments
        for (Region* region : while_region.getRegions()) {
          BlockArgument arg = region->getArgument(output_index);
          LLVM_DEBUG(llvm::dbgs()
                     << "  Propagating device = '" << *device << "' to arg #"
                     << output_index << " of region #"

  if (!main_func_op) return failure();

  SmallVector<StringRef> input_names = GetEntryFunctionInputs(main_func_op);
  SymbolTable symbol_table(module_op);
  SmallVector<AssetFileDef> asset_file_defs;

  for (BlockArgument argument : main_func_op.getArguments()) {
    const int arg_idx = argument.getArgNumber();
    auto asset_op =
        LookupBoundInputOfType<AssetOp>(main_func_op, arg_idx, symbol_table);
    if (!asset_op) continue;

  // args. This will basically be our set increment fn.
  for (Operation* op : non_tpu_ops)
    for (Value input : op->getOperands())
      if (BlockArgument arg = llvm::dyn_cast<BlockArgument>(input))
        // TODO(bfontain): Check that this is actually an argument to the loop
        // body.
        non_tpu_args.insert(arg.getArgNumber());

      std::string value_device = "";

      // If the value is defined as an argument of the func_op, adds it to
      // the argument list of the function that uses this op.
      if (BlockArgument block_arg = mlir::dyn_cast<BlockArgument>(value)) {
        if (StringAttr attr = func_op.getArgAttrOfType<StringAttr>(
                block_arg.getArgNumber(), kTFDeviceAttr)) {
          value_device = attr.getValue().str();
        }

    llvm::SmallDenseMap<llvm::StringRef, TF::ConstOp> const_ops;
    llvm::BitVector arg_indexes_to_remove(func.getNumArguments());

    // Replace arguments with const ops.
    for (BlockArgument argument : func.getArguments()) {
      auto asset = LookupBoundInputOfType<AssetOp>(
          func, argument.getArgNumber(), symbol_table);
      if (asset == nullptr) {
        continue;
      }

    return (arg.hasOneUse() &&
            llvm::isa<quantfork::QuantizeCastOp>(*arg.user_begin()));
  };

  bool need_to_set_input_nodes_quantization_params = false;
  for (const BlockArgument arg : func.getArguments()) {
    auto shaped = mlir::dyn_cast<ShapedType>(arg.getType());
    if (shaped && mlir::isa<FloatType>(shaped.getElementType()) &&
        !has_quantize_op(arg)) {

  // is read in the function body. These read values are just the same argument
  // with type replaced.
  llvm::SmallVector<Value, 4> skipped_args;
  for (auto& it : hoister.GetResources()) {
    BlockArgument arg = mlir::dyn_cast<BlockArgument>(it.first);
    assert(arg && "Expect resources for FuncOp to be its arguments");
    auto type_iter = resource_data_types.find(arg.getArgNumber());
    if (type_iter == resource_data_types.end()) {

  llvm::SmallDenseMap<int64_t, llvm::StringSet<>> result_sets;
  auto insert = [&](Value v, const string& source, const Block& func_block) {
    auto arg = v.dyn_cast<BlockArgument>();
    if (!arg || arg.getOwner() != &func_block) return;
    auto insert_res = result_sets[arg.getArgNumber()].insert(source);
    if (!insert_res.second) return;
    result[arg.getArgNumber()].push_back(source);