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 #"

  // 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();
        }

    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)) {

  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);

    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)) {

  if (auto* resource_op = resource.getDefiningOp()) {
    return resource_op->getAttr(kDeviceAttr);
  } else {
    const auto resource_arg = resource.dyn_cast_or_null<BlockArgument>();
    if (resource_arg && (resource_arg.getOwner() == &(func.front()))) {
      return func.getArgAttrOfType<mlir::StringAttr>(
          resource_arg.getArgNumber(), kFuncDeviceAttr);
    } else {

    if (!arg.hasOneUse() ||
        !mlir::isa<TF::ResourceType>(getElementTypeOrSelf(arg.getType()))) {
      continue;
    }
    auto block_arg = mlir::dyn_cast<BlockArgument>(arg);
    if (!block_arg || block_arg.getOwner() != &replicate.GetBody()) continue;
    assert(replicate_arg_to_execute_arg.count(block_arg.getArgNumber()) == 0 &&

    }
    return success();
  }

 private:
  // Check if op's user or op's user after an identity op is connected to a
  // terminator.
  bool checkIfAnyUserIsConnectedToTermiantor(BlockArgument op) const {
    for (const auto& user : op.getUsers()) {
      if (user->template hasTrait<OpTrait::IsTerminator>()) return true;
      if (auto next_user = dyn_cast_or_null<TF::IdentityOp>(user)) {

      return failure();
    }

    Operation* operand_op = operand.getDefiningOp();
    if (operand_op == nullptr) {
      // When `QuantizeOpT`'s operand does not have a defining op, it means it
      // is a `BlockArgument`. The pattern does not match if there is no op to
      // quantize.
      return failure();
    }

    if (operand_op->hasTrait<OpTrait::ConstantLike>()) {