Add support for unknown dimension in reshape

inducer · Jan 27, 2015 · 84f3624 · 84f3624
1 parent 9c03e9e
commit 84f3624
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/pycuda/gpuarray.py b/pycuda/gpuarray.py
@@ -697,6 +697,15 @@ def reshape(self, *shape):
         if shape == self.shape:
             return self
 
+        if -1 in shape:
+            shape = list(shape)
+            idx = shape.index(-1)
+            size = -reduce(lambda x, y: x * y, shape, 1)
+            shape[idx] = self.size // size
+            if -1 in shape[idx:]:
+                raise ValueError("can only specify one unknown dimension")
+            shape = tuple(shape)
+
         size = reduce(lambda x, y: x * y, shape, 1)
         if size != self.size:
             raise ValueError("total size of new array must be unchanged")

diff --git a/test/test_gpuarray.py b/test/test_gpuarray.py
@@ -758,6 +758,18 @@ def test_reshape(self):
         a_gpu.reshape((4, 32))
         a_gpu.reshape([4, 32])
 
+        # using -1 as unknown dimension
+        assert a_gpu.reshape(-1, 32).shape == (4, 32)
+        assert a_gpu.reshape((32, -1)).shape == (32, 4)
+        assert a_gpu.reshape(((8, -1, 4))).shape == (8, 4, 4)
+
+        throws_exception = False
+        try:
+            a_gpu.reshape(-1, -1, 4)
+        except ValueError:
+            throws_exception = True
+        assert throws_exception
+
     @mark_cuda_test
     def test_view(self):
         a = np.arange(128).reshape(8, 16).astype(np.float32)