Commits

Brian Kearns committed 016e3a6

fix take with multidimensional indices argument

Comments (0)

Files changed (2)

pypy/module/micronumpy/interp_numarray.py

     def take(a, indices, axis, out, mode):
         assert mode == 'raise'
         if axis is None:
-            res = a.ravel()[indices]
+            from numpy import array
+            indices = array(indices)
+            res = a.ravel()[indices.ravel()].reshape(indices.shape)
         else:
             from operator import mul
             if axis < 0: axis += len(a.shape)

pypy/module/micronumpy/test/test_numarray.py

         assert (arange(10).take([1, 2, 1, 1]) == [1, 2, 1, 1]).all()
         raises(IndexError, "arange(3).take([15])")
         a = arange(6).reshape(2, 3)
+        assert a.take(3) == 3
+        assert a.take(3).shape == ()
         assert (a.take([1, 0, 3]) == [1, 0, 3]).all()
+        assert (a.take([[1, 0], [2, 3]]) == [[1, 0], [2, 3]]).all()
         assert (a.take([1], axis=0) == [[3, 4, 5]]).all()
         assert (a.take([1], axis=1) == [[1], [4]]).all()
         assert ((a + a).take([3]) == [6]).all()
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.