Repository: apple/coremltools

Issue:
PyTorch topk not supported after mention in release notes
## 🐞Describe the bug
I am converting a PyTorch model that uses their `torch.topk` operator and I get a runtime error that is is unsupported. Contrary to this, the release [notes](https://github.com/apple/coremltools/releases/tag/4.0b4) for 4.0b4 explicitly say support was added. I have installed coremltools v4.0 which came after that pre-release, so I assume topk should also be supported in my version

## Trace

```
Traceback (most recent call last):
  File "/home/addison/miniconda3/envs/coreml/lib/python3.6/site-packages/IPython/core/interactiveshell.py", line 3331, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-9-21395711530d>", line 1, in <module>
    mlmodel = ct.convert(traced_model, inputs=[ct.ImageType(name='images', shape=dummy.shape)])
  File "/home/addison/miniconda3/envs/coreml/lib/python3.6/site-packages/coremltools/converters/_converters_entry.py", line 183, in convert
    **kwargs
  File "/home/addison/miniconda3/envs/coreml/lib/python3.6/site-packages/coremltools/converters/mil/converter.py", line 129, in mil_convert
    ConverterRegistry, **kwargs)
  File "/home/addison/miniconda3/envs/coreml/lib/python3.6/site-packages/coremltools/converters/mil/converter.py", line 171, in mil_convert_to_proto
    prog = frontend_converter(model, **kwargs)
  File "/home/addison/miniconda3/envs/coreml/lib/python3.6/site-packages/coremltools/converters/mil/converter.py", line 85, in __call__
    return load(*args, **kwargs)
  File "/home/addison/miniconda3/envs/coreml/lib/python3.6/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 83, in load
    raise e
  File "/home/addison/miniconda3/envs/coreml/lib/python3.6/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 75, in load
    prog = converter.convert()
  File "/home/addison/miniconda3/envs/coreml/lib/python3.6/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 224, in convert
    convert_nodes(self.context, self.graph)
  File "/home/addison/miniconda3/envs/coreml/lib/python3.6/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 53, in convert_nodes
    "PyTorch convert function for op '{}' not implemented.".format(node.kind)
RuntimeError: PyTorch convert function for op 'topk' not implemented.
```

## To Reproduce

If `torch.topk` is still not supported, that is okay and I do not need to upload a model and conversion script. I am mainly opening the issue to point out the discrepancy between release notes and actual behavior

## System environment (please complete the following information):
 - coremltools version: 4.0
 - OS: Ubuntu 18.04
 - macOS version (if applicable): NA
 - XCode version (if applicable): NA
 - How you install python: anaconda
 - python version: 3.6.9
 - PyTorch: 1.4.0



Discussion:
yeah that seems to be a gap in testing and a bug. 
The `topk` op was added to the torch `ops.py` but it lacks the decorator `@register_torch_op`, hence it fails to register. 
A temporary workaround is to install coremltools from source and add `@register_torch_op` [here](https://github.com/apple/coremltools/blob/157c1b06997d1f61ea90075e7ad7f5efcb075875/coremltools/converters/mil/frontend/torch/ops.py#L2356)
@aseemw Thank you for the quick workaround! I was able to run the conversion by simply editing the Python code without needing to reinstall. Shall we leave this open until the official release is fixed?

Resolving patch:
diff --git a/coremltools/converters/mil/frontend/torch/ops.py b/coremltools/converters/mil/frontend/torch/ops.py
--- a/coremltools/converters/mil/frontend/torch/ops.py
+++ b/coremltools/converters/mil/frontend/torch/ops.py
@@ -2353,6 +2353,7 @@ def neg(context, node):
     inputs = _get_inputs(context, node, expected=1)
     context.add(mb.mul(x=inputs[0], y=-1, name=node.name))
 
+@register_torch_op
 def topk(context, node):
     inputs = _get_inputs(context, node)
     kwargs = {"name": node.name, "x": inputs[0], "k": inputs[1]}


Test patch:
diff --git a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py b/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py
--- a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py
+++ b/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py
@@ -1432,3 +1432,41 @@ def test_std_4_inputs(self, backend, unbiased, dim, keepdim):
                               kwargs={"unbiased": unbiased, "dim" : dim, "keepdim": keepdim})
         input_shape = (2, 5, 10)
         run_compare_torch(input_shape, model, backend=backend)
+
+class TestTopk:
+    @pytest.mark.parametrize(
+        "backend, largest, shape_dim_k",
+        itertools.product(
+            backends,
+            [True, False],
+            [
+             ((4, 6, 7, 3), -1, 2),
+             ((10, 3, 4), 2, 2),
+             ((10, 5), -2, 3),
+             ((5,), 0, 2)
+             ],
+        ),
+    )
+    def test_topk(self, backend, largest, shape_dim_k):
+        input_shape = shape_dim_k[0]
+        dim = shape_dim_k[1]
+        k = shape_dim_k[2]
+
+        class TopkModel(nn.Module):
+            def __init__(self):
+                super(TopkModel, self).__init__()
+
+            def forward(self, x):
+                return torch.topk(x, k, dim=dim, largest=largest)
+
+        input_data = torch.rand(input_shape)
+        model = TopkModel()
+        expected_results = model(input_data)
+        expected_results = [expected_results.values, expected_results.indices]
+        run_compare_torch(
+            input_data,
+            model,
+            expected_results=expected_results,
+            input_as_shape=False,
+            backend=backend,
+        )
\ No newline at end of file