add aten.topk implementation

[lanluo-nvidia]

Description

Add aten.topk implementation

Type of change

Please delete options that are not relevant and/or add your own.

• New feature (non-breaking change which adds functionality)

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes
• I have added the relevant labels to my PR in so that relevant reviewers are notified

[apbose]

Looks good to me! Pending on CI.

[zewenli98]

@narendasan If TRT doesn't have sorted flag, the outputs of TRT and pytorch might be different. e.g.,

>>> x
tensor([ 1.2096, -0.0950,  0.7965, -0.0960,  0.7188, -1.5562])
>>> torch.ops.aten.topk.default(x, 4, sorted=False)
(tensor([ 0.7188,  1.2096,  0.7965, -0.0950]), tensor([4, 0, 2, 1]))
>>> torch.ops.aten.topk.default(x, 4, sorted=True)
(tensor([ 1.2096,  0.7965,  0.7188, -0.0950]), tensor([0, 2, 4, 1]))

Is the current implementation acceptable?

[apbose]

yeah thats a valid point, and I was also wondering this in the test cases-

In the above you mentioned that topk in TRT returns sorted by default. So was curious if there was a mismatch in the test cases where sorted=False.

Looks like in the test cases, the default flag from torch matches the TRT sorted output.
I missed this point in the next review. @lanluo-nvidia you have more insights?