if (is_func && isa<func::ReturnOp>(user)) {
        assert(!info.IsResultIndexAssigned() &&
               "Expect resource argument to returned no more than once");
        info.result_index = use.getOperandNumber();
        continue;
      }

      auto read = dyn_cast<TF::ReadVariableOp>(user);
      auto write = dyn_cast<TF::AssignVariableOp>(user);
      if (!read && !write) {
        unsupported_use = true;

  for (OpOperand& operand : return_op->getOpOperands()) {
    if (non_tpu_args.contains(operand.getOperandNumber())) {
      if (BlockArgument argument =
              llvm::dyn_cast<BlockArgument>(operand.get())) {
        if (argument.getArgNumber() != operand.getOperandNumber()) {
          return_op.emitOpError()
              << "non TPU ops do not divide state into two pieces.";

  auto return_op = llvm::dyn_cast<mlir::tf_device::ReturnOp>(
      device_cluster.GetBody().getTerminator());

  for (OpOperand& operand : return_op->getOpOperands()) {
    auto operand_idx = operand.getOperandNumber();
    return_operand_map[return_op.getOperand(operand_idx)].push_back(
        operand_idx);
  }
  // `return_value_from_host` and `return_value_from_device` together contain

            // Tensorlist is passed into a while loop, check inside the body
            // function.
            auto inside_uses = while_op.body_function()
                                   .getArgument(use.getOperandNumber())
                                   .getUses();
            for (auto &inside_use : llvm::make_early_inc_range(inside_uses)) {
              if (TF::TensorListSetItemOp set_op =

      auto affine_user = llvm::dyn_cast<mlir::AffineQuantizedOpInterface>(user);
      if (affine_user &&
          affine_user.GetAffineOperandIndex() == use.getOperandNumber() &&
          affine_user.RequiredNarrowRangeAffineOperand())
        continue;
      // Create a fully range quantized constant.
      if (full_range_const == value) {
        mlir::quant::QuantizedType new_qtype;