Added test with 3 values of k

tests/transformation/test_topk_insert.py

+import onnx
+from finn.util.test import get_test_model_trained
+import brevitas.onnx as bo
+import numpy as np
+import onnx.numpy_helper as nph
+import torch
+
+from finn.core.modelwrapper import ModelWrapper
+from finn.transformation.general import GiveReadableTensorNames, GiveUniqueNodeNames
+from finn.transformation.infer_shapes import InferShapes
+from finn.transformation.infer_datatypes import InferDataTypes
+from finn.transformation.fold_constants import FoldConstants
+from finn.transformation.insert_topk import InsertTopK
+
+import finn.core.onnx_exec as oxe
+from pkgutil import get_data
+
+import pytest
+
+export_onnx_path = "test_output_lfc.onnx"
+
+@pytest.mark.parametrize("k", [1,5,10])
+def test_topk_insert(k):
+    tfc = get_test_model_trained("TFC", 1, 1)
+    bo.export_finn_onnx(tfc, (1, 1, 28, 28), export_onnx_path)
+    model = ModelWrapper(export_onnx_path)
+
+    #do transformations (no topk)
+    model = model.transform(InferShapes())
+    model = model.transform(FoldConstants())
+    model = model.transform(GiveUniqueNodeNames())
+    model = model.transform(GiveReadableTensorNames())
+    model = model.transform(InferDataTypes())
+
+    #verification: generate random input, run through net, streamline, run again, check that output is top-k
+    raw_i = get_data("finn", "data/onnx/mnist-conv/test_data_set_0/input_0.pb")
+    input_tensor = onnx.load_tensor_from_string(raw_i)
+    input_brevitas = torch.from_numpy(nph.to_array(input_tensor)).float()
+    output_golden = tfc.forward(input_brevitas).detach().numpy()
+    output_golden_topk = np.flip(output_golden.flatten().argsort())[:k]
+
+    input_dict = {"global_in": nph.to_array(input_tensor)}
+    output_dict = oxe.execute_onnx(model, input_dict)
+    output_pysim = output_dict[list(output_dict.keys())[0]]
+
+    #insert top-k
+    model = model.transform(InsertTopK(k))
+    model = model.transform(GiveUniqueNodeNames())
+    model = model.transform(GiveReadableTensorNames())
+    model = model.transform(InferShapes())
+
+    #verify output of top-k
+    output_dict_topk = oxe.execute_onnx(model, input_dict)
+    output_pysim_topk = output_dict_topk[list(output_dict_topk.keys())[0]]
+
+    assert np.array_equal(output_golden_topk.flatten(), output_pysim_topk.astype(np.int).flatten())
+