Missing TopK operator for Onnx

sonos / tract

Tiny, no-nonsense, self-contained, Tensorflow and ONNX inference

Other

2.24k stars 214 forks source link

Missing TopK operator for Onnx #974

Closed tgolsson closed 1 year ago

tgolsson commented 1 year ago

Hello!

The operator TopK (https://onnx.ai/onnx/operators/onnx__TopK.html) does not seem to be implemented. This is used as a sparse activation function in some architectures.

Error: Failed analyse for node #88 "model_4/kwta_layer/PartitionedCall/TopKV2" Unimplemented(TopK)

I assume this'd need to be implemented both for ONNX plus an NNEF-tract extension as NNEF does not have a native TopK as I understand it.

kali commented 1 year ago

Hey @tgolsson... can I assume that K will be a constant known at loading time (in which case it's pretty easy) or will models make k variable in real life ?

tgolsson commented 1 year ago

I think the spec says dynamic but for our use cases fixed is fine.