auto begin_type = mlir::dyn_cast<RankedTensorType>(begin_indices.getType());
    auto size_type = mlir::dyn_cast<RankedTensorType>(sizes.getType());

    if (!input_ty || !begin_type || !size_type ||
        !begin_type.hasStaticShape() || !size_type.hasStaticShape() ||
        begin_type.getRank() != 1 || size_type.getRank() != 1) {
      return failure();
    }

    func::FuncOp ret = queue_.front();
    queue_.pop();
    queue_set_.erase(ret);
    return ret;
  }

  // Returns the current size of the queue.
  std::queue<func::FuncOp>::size_type QueueSize() const {
    return queue_.size();
  }

  Dialect* const tf_dialect_;

 private:
  // Returns whether the result of an operation could be updated to a new