Add negative dimension to transpose and tests (#792)

goelhardik · soumith · commit c93c884ee2f4 · 2017-03-03T09:31:22.000-05:00
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -7,7 +7,7 @@
 import tempfile
 import unittest
 import warnings
-from itertools import product
+from itertools import product, combinations
 from common import TestCase, iter_indices, TEST_NUMPY, run_tests, download_file, skipIfNoLapack
 
 if TEST_NUMPY:
@@ -2987,6 +2987,23 @@ def test_Size(self):
         self.assertIsInstance(x[:-1], torch.Size)
         self.assertIsInstance(x + x, torch.Size)
 
+    def test_transpose_neg(self):
+        x = torch.randn(10, 20, 30)
+        ndim = 3
+
+        for i, j in combinations(range(ndim), 2):
+            a = x.transpose(i, j)
+            b = x.transpose(i - ndim, j - ndim)
+            self.assertEqual(a, b)
+
+            a = torch.transpose(x, i, j)
+            b = torch.transpose(x, i - ndim, j - ndim)
+            self.assertEqual(a, b)
+
+            a = x.clone()
+            x.transpose_(i, j)
+            x.transpose_(i - ndim, j - ndim)
+            self.assertEqual(a, x)
 
 if __name__ == '__main__':
     run_tests()
diff --git a/torch/csrc/generic/methods/Tensor.cwrap b/torch/csrc/generic/methods/Tensor.cwrap
@@ -342,21 +342,41 @@ PyObject * THPTensor_(stride)(PyObject *self, PyObject *args, PyObject *kwargs)
     - THBoolTensor* mask
 ]]
 
+#if IS_CUDA
+THTensor* THTensor_(transpose_neg)(THCState* state, THTensor *self, THTensor *src, int dim0, int dim1)
+#else
+THTensor* THTensor_(transpose_neg)(THTensor *self, THTensor *src, int dim0, int dim1)
+#endif
+{
+  int ndim = self->nDimension;
+  if (dim0 < 0)
+    dim0 += ndim;
+  if (dim1 < 0)
+    dim1 += ndim;
+  if (src != NULL) {
+    THTensor_(transpose)(LIBRARY_STATE self, src, dim0, dim1);
+    return NULL;
+  } else {
+    return THTensor_(newTranspose)(LIBRARY_STATE self, dim0, dim1);
+  }
+}
+
 [[
   name: transpose
   with_stateless: True
-  cname: newTranspose
+  cname: transpose_neg
   cpu_half: True
   return: THTensor*
   arguments:
     - THTensor* self
+    - CONSTANT NULL
     - long dim0
     - long dim1
 ]]
 
 [[
   name: transpose_
-  cname: transpose
+  cname: transpose_neg
   cpu_half: True
   return: self
   arguments:
