Repository: pytorch/pytorch

Issue:
[feature request] torch.argmax / torch.argmin
To save indexing into output of torch.max / torch.min and make code more readable. NumPy / TensorFlow also has it.


Discussion:
How is this feature going? Has it been added?
@thecortex note that you can currently obtain the argmax as the second return value of `torch.max` when a dimension is specified. Same for argmin

Resolving patch:
diff --git a/tools/autograd/gen_python_functions.py b/tools/autograd/gen_python_functions.py
--- a/tools/autograd/gen_python_functions.py
+++ b/tools/autograd/gen_python_functions.py
@@ -18,7 +18,7 @@
     '.*_backward', '.*_backward_(out|input|weight|bias)', '.*_forward',
     '.*_forward_out', 'sparse_raw_resize_', '_unsafe_view', 'tensor',
     'sparse_coo_tensor', '_arange.*', '_range.*', '_linspace.*', '_logspace.*',
-    '_indexCopy_',
+    '_indexCopy_', 'max_values', 'min_values', 'argmax', 'argmin'
 ]
 
 PY_VARIABLE_METHODS_CPP = CodeTemplate.from_file(template_path + '/python_variable_methods.cpp')
diff --git a/torch/functional.py b/torch/functional.py
--- a/torch/functional.py
+++ b/torch/functional.py
@@ -4,6 +4,8 @@
 import math
 
 __all__ = [
+    'argmax',
+    'argmin',
     'bartlett_window',
     'btrifact',
     'btriunpack',
@@ -378,3 +380,75 @@ def unique(input, sorted=False, return_inverse=False):
         return output, inverse_indices
     else:
         return output
+
+
+def argmax(input, dim=None, keepdim=False):
+    """Returns the indices of the maximum values of a tensor across a dimension.
+
+    This is the second value returned by :meth:`torch.max`. See its
+    documentation for the exact semantics of this method.
+
+    Args:
+        input (Tensor): the input tensor
+        dim (int): the dimension to reduce. If ``None``, the argmax of the
+            flattened input is returned.
+        keepdim (bool): whether the output tensors have :attr:`dim`
+            retained or not. Ignored if ``dim=None``.
+
+    Example::
+
+        >>> a = torch.randn(4, 4)
+        >>> a
+
+         2.3461  0.0056  1.4846  0.3911
+        -1.3584 -1.0066  0.0530  1.1754
+        -0.7929 -0.3194 -1.4865  0.4020
+         0.1101  0.6694  1.3456  0.8235
+        [torch.FloatTensor of size (4,4)]
+
+        >>> torch.argmax(a, dim=1)
+        0
+        3
+        3
+        2
+        [torch.LongTensor of size (4,)]
+    """
+    if dim is None:
+        return torch._argmax(input.contiguous().view(-1), dim=0, keepdim=False)
+    return torch._argmax(input, dim, keepdim)
+
+
+def argmin(input, dim=None, keepdim=False):
+    """Returns the indices of the minimum values of a tensor across a dimension.
+
+    This is the second value returned by :meth:`torch.min`. See its
+    documentation for the exact semantics of this method.
+
+    Args:
+        input (Tensor): the input tensor
+        dim (int): the dimension to reduce. If ``None``, the argmin of the
+            flattened input is returned.
+        keepdim (bool): whether the output tensors have :attr:`dim`
+            retained or not. Ignored if ``dim=None``.
+
+    Example::
+
+        >>> a = torch.randn(4, 4)
+        >>> a
+
+         2.3461  0.0056  1.4846  0.3911
+        -1.3584 -1.0066  0.0530  1.1754
+        -0.7929 -0.3194 -1.4865  0.4020
+         0.1101  0.6694  1.3456  0.8235
+        [torch.FloatTensor of size (4,4)]
+
+        >>> torch.argmin(a, dim=1)
+         1
+         0
+         2
+         0
+        [torch.LongTensor of size (4,)]
+    """
+    if dim is None:
+        return torch._argmin(input.contiguous().view(-1), dim=0, keepdim=False)
+    return torch._argmin(input, dim, keepdim)
diff --git a/torch/tensor.py b/torch/tensor.py
--- a/torch/tensor.py
+++ b/torch/tensor.py
@@ -214,6 +214,14 @@ def share_memory_(self):
     def view_as(self, tensor):
         return self.view(tensor.size())
 
+    def argmax(self, dim=None, keepdim=False):
+        r"""See :func:`torch.argmax`"""
+        return torch.argmax(self, dim, keepdim)
+
+    def argmin(self, dim=None, keepdim=False):
+        r"""See :func:`torch.argmin`"""
+        return torch.argmin(self, dim, keepdim)
+
     def btrifact(self, info=None, pivot=True):
         r"""See :func:`torch.btrifact`
         """


Test patch:
diff --git a/test/test_torch.py b/test/test_torch.py
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -574,6 +574,28 @@ def _test_dim_reduction(self, cast):
             self.assertEqual(x.min(0), (torch.FloatTensor([-1, 2, 1]), torch.FloatTensor([0, 0, 0])))
             self.assertEqual(x.min(1), (torch.FloatTensor([-1, 3]), torch.FloatTensor([0, 1])))
 
+        for dtype in types:
+            if dtype == torch.uint8:  # Doesn't support negative values
+                continue
+            x = cast(torch.tensor(example, dtype=dtype))
+            self.assertEqual(x.argmax().item(), 5)
+            self.assertEqual(x.argmax(dim=0), torch.FloatTensor([1, 1, 1]))
+            self.assertEqual(x.argmax(dim=1), torch.FloatTensor([1, 2]))
+            self.assertEqual(x.argmax(dim=0, keepdim=True), torch.FloatTensor([[1, 1, 1]]))
+            # test that non-contiguous tensors work
+            self.assertEqual(x[:, :2].argmax().item(), 2)
+
+        for dtype in types:
+            if dtype == torch.uint8:  # Doesn't support negative values
+                continue
+            x = cast(torch.tensor(example, dtype=dtype))
+            self.assertEqual(x.argmin().item(), 0)
+            self.assertEqual(x.argmin(dim=0), torch.FloatTensor([0, 0, 0]))
+            self.assertEqual(x.argmin(dim=1), torch.FloatTensor([0, 1]))
+            self.assertEqual(x.argmin(dim=1, keepdim=True), torch.FloatTensor([[0], [1]]))
+            # test that non-contiguous tensors work
+            self.assertEqual(x[:, :2].argmin().item(), 0)
+
         dim_red_fns = [
             "mean", "median", "mode", "norm", "prod",
             "std", "sum", "var", "max", "min"]