Add support for topk (jonasrauber#24)

* Add support for topk

* Fix return type

* reformat with black

* style tweaks: tuple unpacking

* Testcase that is not already sorted

* Use integers for more determinism

docs/api/tensor.md

@@ -147,6 +147,11 @@ Tensor.argmax(self:~TensorType, axis:Union[int, NoneType]=None) -> ~TensorType
 Tensor.argsort(self:~TensorType, axis:int=-1) -> ~TensorType
 ```
 
+## topk
+```python
+Tensor.topk(self:~TensorType, k:int, sorted:bool=True) -> Tuple[~TensorType, ~TensorType]
+```
+
 ## uniform
 ```python
 Tensor.uniform(self:~TensorType, shape:Union[Tuple[int, ...], int], low:float=0.0, high:float=1.0) -> ~TensorType

eagerpy/framework.py

@@ -122,6 +122,10 @@ def sort(t: TensorType, axis: int = -1) -> TensorType:
     return t.sort(axis=axis)
 
 
+def topk(t: TensorType, k: int, sorted: bool = True) -> Tuple[TensorType, TensorType]:
+    return t.topk(k, sorted=sorted)
+
+
 def uniform(
     t: TensorType, shape: ShapeOrScalar, low: float = 0.0, high: float = 1.0
 ) -> TensorType:

eagerpy/tensor/jax.py

@@ -174,6 +174,20 @@ def argsort(self: TensorType, axis: int = -1) -> TensorType:
     def sort(self: TensorType, axis: int = -1) -> TensorType:
         return type(self)(self.raw.sort(axis=axis))
 
+    def topk(
+        self: TensorType, k: int, sorted: bool = True
+    ) -> Tuple[TensorType, TensorType]:
+        # argpartition not yet implemented
+        # wrapping indexing not supported in take()
+        n = self.raw.shape[-1]
+        idx = np.take(np.argsort(self.raw), np.arange(n - k, n), axis=-1)
+        val = np.take_along_axis(self.raw, idx, axis=-1)
+        if sorted:
+            perm = np.flip(np.argsort(val, axis=-1), axis=-1)
+            idx = np.take_along_axis(idx, perm, axis=-1)
+            val = np.take_along_axis(self.raw, idx, axis=-1)
+        return type(self)(val), type(self)(idx)
+
     def uniform(
         self: TensorType, shape: ShapeOrScalar, low: float = 0.0, high: float = 1.0
     ) -> TensorType:

eagerpy/tensor/numpy.py

@@ -125,6 +125,17 @@ def argsort(self: TensorType, axis: int = -1) -> TensorType:
     def sort(self: TensorType, axis: int = -1) -> TensorType:
         return type(self)(np.sort(self.raw, axis=axis))
 
+    def topk(
+        self: TensorType, k: int, sorted: bool = True
+    ) -> Tuple[TensorType, TensorType]:
+        idx = np.take(np.argpartition(self.raw, k - 1), np.arange(-k, 0), axis=-1)
+        val = np.take_along_axis(self.raw, idx, axis=-1)
+        if sorted:
+            perm = np.flip(np.argsort(val, axis=-1), axis=-1)
+            idx = np.take_along_axis(idx, perm, axis=-1)
+            val = np.take_along_axis(self.raw, idx, axis=-1)
+        return type(self)(val), type(self)(idx)
+
     def uniform(
         self: TensorType, shape: ShapeOrScalar, low: float = 0.0, high: float = 1.0
     ) -> TensorType:

eagerpy/tensor/pytorch.py

@@ -203,6 +203,12 @@ def argsort(self: TensorType, axis: int = -1) -> TensorType:
     def sort(self: TensorType, axis: int = -1) -> TensorType:
         return type(self)(self.raw.sort(dim=axis).values)  # type: ignore
 
+    def topk(
+        self: TensorType, k: int, sorted: bool = True
+    ) -> Tuple[TensorType, TensorType]:
+        values, indices = self.raw.topk(k, sorted=sorted)
+        return type(self)(values), type(self)(indices)
+
     def uniform(
         self: TensorType, shape: ShapeOrScalar, low: float = 0.0, high: float = 1.0
     ) -> TensorType:

eagerpy/tensor/tensor.py

@@ -309,6 +309,12 @@ def argsort(self: TensorType, axis: int = -1) -> TensorType:
     def sort(self: TensorType, axis: int = -1) -> TensorType:
         ...
 
+    @abstractmethod
+    def topk(
+        self: TensorType, k: int, sorted: bool = True
+    ) -> Tuple[TensorType, TensorType]:
+        ...
+
     @abstractmethod
     def uniform(
         self: TensorType, shape: ShapeOrScalar, low: float = 0.0, high: float = 1.0

eagerpy/tensor/tensorflow.py

@@ -176,6 +176,12 @@ def argsort(self: TensorType, axis: Optional[int] = -1) -> TensorType:
     def sort(self: TensorType, axis: Optional[int] = -1) -> TensorType:
         return type(self)(tf.sort(self.raw, axis=axis))
 
+    def topk(
+        self: TensorType, k: int, sorted: bool = True
+    ) -> Tuple[TensorType, TensorType]:
+        values, indices = tf.math.top_k(self.raw, k, sorted=sorted)
+        return type(self)(values), type(self)(indices)
+
     @samedevice
     def uniform(
         self: TensorType, shape: ShapeOrScalar, low: float = 0.0, high: float = 1.0

tests/test_main.py

@@ -1206,6 +1206,20 @@ def test_sort(dummy: Tensor) -> Tensor:
     return ep.sort(t)
 
 
+@compare_all
+def test_topk_values(dummy: Tensor) -> Tensor:
+    t = (ep.arange(dummy, 27).reshape((3, 3, 3)) ** 2 * 10000 % 1234).float32()
+    values, _ = ep.topk(t, 2)
+    return values
+
+
+@compare_all
+def test_topk_indices(dummy: Tensor) -> Tensor:
+    t = -(ep.arange(dummy, 27).reshape((3, 3, 3)) ** 2 * 10000 % 1234).float32()
+    _, indices = ep.topk(t, 2)
+    return indices
+
+
 @compare_all
 def test_transpose(dummy: Tensor) -> Tensor:
     t = ep.arange(dummy, 8).float32().reshape((2, 4))