upd (#1113)

6beef85b · Zihao Ye · GitHub · 41114b87 · 6beef85b
Unverified Commit 6beef85b authored Dec 17, 2019 by Zihao Ye Committed by GitHub Dec 17, 2019
Show whitespace changes
Inline Side-by-side

Showing with 8 additions and 0 deletions

python/dgl/backend/numpy/tensor.py python/dgl/backend/numpy/tensor.py +8 -0

No files found.
--- a/python/dgl/backend/numpy/tensor.py
+++ b/python/dgl/backend/numpy/tensor.py
@@ -110,6 +110,14 @@ def argsort(input, dim, descending):
        return np.argsort(-input, axis=dim)
    return np.argsort(input, axis=dim)

+def topk(input, k, dim, descending=True):
+    topk_indices = argtopk(input, k, dim, descending)
+    return np.take_along_axis(input, topk_indices, axis=dim)
+
+def argtopk(input, k, dim, descending=True):
+    sort_indces = argsort(input, dim, descending)
+    return slice_axis(sort_indces, dim, 0, k)
+
 def exp(input):
    return np.exp(input)