Add reversed(torch.Tensor) (#9216)

vishwakftw · facebook-github-bot · commit cd3e067e4654 · 2018-07-12T19:42:07.000-07:00
Summary: Closes #3376 Pull Request resolved: #9216 Differential Revision: D8753933 Pulled By: soumith fbshipit-source-id: 5dac9b8b11ff34a205b6478db99b02fda8bd9cce
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -6361,6 +6361,16 @@ def _test_flip(self, use_cuda=False):
     def test_flip(self):
         self._test_flip(self, use_cuda=False)
 
+    def test_reversed(self):
+        val = torch.arange(0, 10)
+        self.assertEqual(reversed(val), torch.arange(9, -1, -1))
+
+        val = torch.arange(1, 10).view(3, 3)
+        self.assertEqual(reversed(val), torch.tensor([[7, 8, 9], [4, 5, 6], [1, 2, 3]]))
+
+        val = torch.tensor(42)
+        self.assertEqual(reversed(val), torch.tensor(42))
+
     def test_storage(self):
         v = torch.randn(3, 5)
         self.assertEqual(v.storage()[0], v.data[0][0])
diff --git a/torch/tensor.py b/torch/tensor.py
@@ -231,6 +231,13 @@ def view_as(self, tensor):
         """
         return self.view(tensor.size())
 
+    def __reversed__(self):
+        r"""Reverses the tensor along dimension 0."""
+        if self.dim() == 0:
+            return self
+        else:
+            return self.flip(0)
+
     def argmax(self, dim=None, keepdim=False):
         r"""See :func:`torch.argmax`"""
         return torch.argmax(self, dim, keepdim)
