Add meta impl for topk

anjali411 · anjali411 · commit 8f143f07dc6f · 2022-11-08T20:09:57.000Z
ghstack-source-id: 4044cf2 Pull Request resolved: #88694
diff --git a/test/test_proxy_tensor.py b/test/test_proxy_tensor.py
@@ -1352,7 +1352,6 @@ def f(a, b, c, d, e):
     xfail('take_along_dim', ''),  # dtype of indices should be Long but got Float
     xfail('take', ''),  # aten.take.default - couldn't find symbolic meta function/decomposition
     xfail('tensordot', ''),  # aten.size.default - couldn't find symbolic meta function/decomposition
-    xfail('topk', ''),  # aten.topk.default - couldn't find symbolic meta function/decomposition
     xfail('trapz', ''),  # aten.size.default - couldn't find symbolic meta function/decomposition
     xfail('trapezoid', ''),  # aten.size.default - couldn't find symbolic meta function/decomposition
     xfail('triangular_solve', ''),  # aten.triangular_solve.default - couldn't find symbolic meta function/decomposition
diff --git a/torch/_meta_registrations.py b/torch/_meta_registrations.py
@@ -1685,6 +1685,21 @@ def scalar_tensor(s, dtype=None, layout=None, device=None, pin_memory=None):
         (), dtype=dtype, layout=layout, device=device, pin_memory=pin_memory
     )
 
+@register_meta(aten.topk.default)
+def topk_meta(self, k, dim=-1, largest=True, sorted=True):
+    # From aten/src/ATen/native/Sorting.cpp
+    dim = maybe_wrap_dim(dim, self.dim(), wrap_scalar=True)
+    check(
+        k >= 0 and k <= (self.size(dim) if self.dim() > 0 else 1),
+        lambda: "selected index k out of range",
+    )
+    sliceSize = 1 if self.dim() == 0 else self.size(dim)
+    check(k >= 0 and k <= sliceSize, lambda: "k not in range for dimension")
+
+    topKSize = list(self.shape)
+    if len(topKSize) > 0:
+        topKSize[dim] = k
+    return self.new_empty(topKSize), self.new_empty(topKSize, dtype=torch.int64)
 
 # We must also trigger meta registrations from PrimTorch ref
 # decompositions
