Repository: opencv/opencv
File: modules/dnn/src/layers/topk2_layer.cpp (C++)

Code before:
    TopK2LayerImpl(const LayerParams& params)
    {
        setParamsFrom(params);
        axis    = params.get<int>("axis", -1);
        largest = params.get<int>("largest", 1) == 1;
        sorted  = params.get<int>("sorted", 1) == 1;
        CV_CheckTrue(sorted, "TopK2: sorted == false is not supported");
        if (params.has("k")) {
            K = params.get<int>("k");
            CV_CheckGT(K, 0, "TopK2: K needs to be a positive integer");
            dynamicK = false;
        } else {
            dynamicK = true;
            K = 0;
        }
    }

    bool supportBackend(int backendId) CV_OVERRIDE
    {
        return backendId == DNN_BACKEND_OPENCV;
    }

    bool getMemoryShapes(const std::vector<MatShape>& inputs,
                         const int requiredOutputs,
                         std::vector<MatShape>& outputs,
                         std::vector<MatShape>& internals) const CV_OVERRIDE
    {
        CV_Assert(!inputs.empty());
        const MatShape& inShape = inputs[0];
        int inputDims = static_cast<int>(inShape.size());
        int a = normalize_axis(axis, inputDims);
        CV_Assert(a >= 0 && a < inputDims);

        int outK = dynamicK ? inShape[a] : K;
        CV_CheckLT(outK, inShape[a] + 1, "TopK2: K is out of range");

        MatShape outShape = inShape;
        outShape[a] = outK;
        outputs.assign(2, outShape);
        internals.clear();
        return false;
    }

    void getTypes(const std::vector<MatType>& inputs,
                  const int requiredOutputs,
                  const int requiredInternals,
                  std::vector<MatType>& outputs,
                  std::vector<MatType>& internals) const CV_OVERRIDE
    {
        outputs.resize(2);
        outputs[0] = inputs.front();

Review comment (bug):
I believe, it's wrong. In the case of dynamic TopK "K" parameter should be taken from the tensor itself, not shape of the tensor.

Code after:
    }

    bool getMemoryShapes(const std::vector<MatShape>& inputs,
                         const int requiredOutputs,
                         std::vector<MatShape>& outputs,
                         std::vector<MatShape>& internals) const CV_OVERRIDE
    {
        CV_Assert(!inputs.empty());
        CV_Assert(K > 0);
        const MatShape& inShape = inputs[0];
        int inputDims = static_cast<int>(inShape.size());
        int a = normalize_axis(axis, inputDims);
        CV_Assert(a >= 0 && a < inputDims);
        CV_CheckLT(K, inShape[a] + 1, "TopK2: K is out of range");

        MatShape outShape = inShape;
        outShape[a] = K;
        outputs.assign(2, outShape);
        internals.clear();
        return false;
    }

    void getTypes(const std::vector<MatType>& inputs,
                  const int requiredOutputs,
                  const int requiredInternals,
                  std::vector<MatType>& outputs,
                  std::vector<MatType>& internals) const CV_OVERRIDE
    {
        outputs.resize(2);
        outputs[0] = inputs.front();
        outputs[1] = CV_64S;
    }

private:
    template<typename T, typename WT = T>
    void FindTopK(const Mat &input, Mat &output_vals, Mat &output_idxs, int normalized_axis, int kVal)
    {
        const auto inShape = shape(input);
        size_t outer = std::accumulate(inShape.begin(), inShape.begin() + normalized_axis, 1, std::multiplies<int>());
        size_t inner = std::accumulate(inShape.begin() + normalized_axis + 1, inShape.end(), 1, std::multiplies<int>());
        int dimAxis = inShape[normalized_axis];

        auto worker = [&](const Range &r) {
            const T* inPtr = input.ptr<T>();
            T* valPtr = output_vals.ptr<T>();
            int64_t* idxPtr = output_idxs.ptr<int64_t>();

            std::vector<std::pair<uint32_t, WT>> sortbuf(dimAxis);
            for (size_t b = r.start; b < r.end; ++b) {
                for (size_t j = 0; j < inner; ++j) {
                    size_t offset = b * dimAxis * inner + j;