PROCESSOR-SDK-J784S4: Unexpected indices shape from TIDL TopK: expected (1,1,30,1), got (1,4,30,1)

Hi Ghassen,

It appears to be an output problem.  I have attached a Jupyter notebook with with a tensor shape of [1,1,10,1] so it is easier to comprehend (100 is a lot of numbers to go through).  I also attached the script to generate a topk model with [1,1,10,1] tensors.  The only point to bring up about the notebook is you need to have a working 10_01_04_00 environment setup. 

Set your 10_01_04_00 environment  in modeldir = '<your_path>'+VERSION+'/edgeai-tidl-tools/tools/AM69A/tidl_tools/'    (cell 4).

You must read the output .bin file as 32 bit values, the data is in float32 while the the indexes are in int32.

index=np.fromfile("out/my_data.bin",dtype=np.uint32)
data=np.fromfile("out/my_data.bin",dtype=np.float32)

If you do not like to use Jupyter notebooks, just copy the -- parameters from the import and inference cell and run from an import/inference config file.

Regards,

Chris

import torch
import torch.nn as nn

class SimpleTopKModel(nn.Module):
    def forward(self, x):
        # Returns the top 3 largest elements along the last dimension
        values, indices = torch.topk(x, k=3, dim=2,largest=True,sorted=True)
        return values, indices

# Create an instance of the model
model = SimpleTopKModel()

# Create a sample input tensor
dummy_input = torch.randn(1,1,10,1) # Example: batch size 1, 10 elements

# Export the model to ONNX
torch.onnx.export(model,
                  dummy_input,
                  "topk_model.onnx",
                  input_names=['input'],
                  output_names=['values', 'indices'],
                  opset_version=11) # Specify an appropriate opset version

print("Model exported to topk_model.onnx")