Merge pull request #8643 from eric-wieser/fix-8642

eric-wieser · web-flow · commit f3682860cfe2 · 2017-04-11T17:42:59.000+01:00
BUG: Fix double-wrapping of object scalars
diff --git a/numpy/lib/shape_base.py b/numpy/lib/shape_base.py
@@ -103,8 +103,10 @@ def apply_along_axis(func1d, axis, arr, *args, **kwargs):
     in_dims = list(range(nd))
     inarr_view = transpose(arr, in_dims[:axis] + in_dims[axis+1:] + [axis])
 
-    # compute indices for the iteration axes
+    # compute indices for the iteration axes, and append a trailing ellipsis to
+    # prevent 0d arrays decaying to scalars, which fixes gh-8642
     inds = ndindex(inarr_view.shape[:-1])
+    inds = (ind + (Ellipsis,) for ind in inds)
 
     # invoke the function on the first item
     try:
diff --git a/numpy/lib/tests/test_shape_base.py b/numpy/lib/tests/test_shape_base.py
@@ -159,6 +159,21 @@ def empty_to_1(x):
         assert_equal(actual, np.ones(10))
         assert_raises(ValueError, np.apply_along_axis, empty_to_1, 0, a)
 
+    def test_with_iterable_object(self):
+        # from issue 5248
+        d = np.array([
+            [set([1, 11]), set([2, 22]), set([3, 33])],
+            [set([4, 44]), set([5, 55]), set([6, 66])]
+        ])
+        actual = np.apply_along_axis(lambda a: set.union(*a), 0, d)
+        expected = np.array([{1, 11, 4, 44}, {2, 22, 5, 55}, {3, 33, 6, 66}])
+
+        assert_equal(actual, expected)
+
+        # issue 8642 - assert_equal doesn't detect this!
+        for i in np.ndindex(actual.shape):
+            assert_equal(type(actual[i]), type(expected[i]))
+
 
 class TestApplyOverAxes(TestCase):
     def test_simple(self):
