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

    // created entry block arguments in the function body.
    {
      Block &entry_block = outlined_func.getBody().front();
      auto loc = outlined_func.getLoc();
      for (Value operand : operands) {
        BlockArgument newArg = entry_block.addArgument(operand.getType(), loc);
        replaceAllUsesInRegionWith(operand, newArg, outlined_func.getBody());
      }
    }

  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 {

  }

  TF::ConstOp tensor_names_const =
      CreateTensorNamesConst(tensor_names, builder);
  TF::ConstOp shape_and_slices_const =
      CreateShapeAndSlicesConst(tensor_names.size(), builder);

  BlockArgument filename_arg = save_func.getArgument(0);
  return builder.create<TF::SaveV2Op>(
      NameLoc::get(builder.getStringAttr(kTfQuantSaveV2OpName)),
      /*prefix=*/filename_arg, tensor_names_const, shape_and_slices_const,

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

// if the given operation is the only user.
Operation* GetInputOpWithOneUse(Operation* op, int opr_num) {
  if (opr_num >= op->getNumOperands()) return nullptr;
  auto opr = op->getOperand(opr_num);
  if (llvm::isa<BlockArgument>(opr)) return nullptr;
  auto* res = opr.getDefiningOp();
  if (!res->hasOneUse()) return nullptr;
  return res;
}

// Checks if the given operand of given operation refers to a splat constant
// with given val.

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

  return new_device;
}

// Propagates devices from function arguments.
void PropagateDevicesFromArguments(
    func::FuncOp func,
    llvm::DenseMap<Value, llvm::StringRef>& value_to_device) {
  for (BlockArgument& arg : func.getArguments()) {
    auto arg_device_attr =
        func.getArgAttrOfType<StringAttr>(arg.getArgNumber(), kFuncDeviceAttr);
    if (!arg_device_attr || arg_device_attr.getValue().empty() ||

    constexpr char kShardingAttr[] = "mhlo.sharding";
    if (auto sharding =
            func.getArgAttrOfType<mlir::StringAttr>(i, kShardingAttr)) {
      if (!sharding.getValue().empty()) {
        BlockArgument arg = func.getArgument(i);
        // TODO(hinsu): Instead of setting both 'sharding' and '_XlaSharding'
        // attributes, only set the 'sharding' attribute. Both attributes are