BUG: Fix crash on 0d return value in apply_along_axis #8441

eric-wieser · 2017-01-02T23:06:59Z

Also:
ENH: Support arbitrary dimensionality of return value
MAINT: remove special casing

Now supports

>>> data = np.arange(6).reshape(2, 3)
>>> data
array([[0, 1, 2],
       [3, 4, 5]])
>>> np.apply_along_axis(np.diag, -1, data)  #would previously crash
array([[[0, 0, 0],
        [0, 1, 0],
        [0, 0, 2]],

       [[3, 0, 0],
        [0, 4, 0],
        [0, 0, 5]]])
>>> np.apply_along_axis(lambda x: np.array(np.sum(x)), 0, data)  #would previously crash
array([3, 5, 7])

Addresses discussion from #8363

eric-wieser · 2017-01-02T23:20:27Z

numpy/lib/shape_base.py

+    # save the first result
+    outarr_view[tuple(ind)] = res
+    k = 1
+    while k < Ntot:


Is this really likely to be faster than using ndindex here?

eric-wieser · 2017-01-02T23:40:55Z

numpy/lib/shape_base.py

+    inarr_view = transpose(arr, dims_in[:axis] + dims_in[axis+1:] + [axis])
+
+    # indices
+    inds = ndindex(inarr_view.shape[:nd-1])


To be honest, most of this function could probably be replaced with a carefully constructed nditer that does the transposes for us, but that's probably not going to increase performance much anyway

eric-wieser · 2017-01-03T00:23:51Z

Right now, this happens:

>>> evil = lambda x: 1 if x.all() else [2, 3]
# silent broadcasting
>>> apply_along_axis(evil, 0, np.array([[False, True, True]]))))
array([[2, 1, 1],
       [3, 1, 1]])
# same function, different order
>>> apply_along_axis(evil, 0, np.array([[True, False, True]]))))
ValueError: setting an array element with a sequence.

Would it be better if both raised ValueError: shape must be the same on each invocation, at the expense of a check on every iteration?

eric-wieser · 2017-01-03T01:44:09Z

numpy/lib/shape_base.py

+    if not isinstance(res, matrix):
+        outarr = res.__array_prepare__(outarr)
+    if outshape != outarr.shape:
+        raise ValueError('__array_prepare__ should not change the shape of the resultant array')


Is this assumption (__array_prepare__(x).shape == x.shape) supposed to be valid? How do we deal with matrix violating it elsewhere?

rgommers · 2017-01-03T02:43:19Z

numpy/lib/tests/test_shape_base.py

+        )
+
+        result = apply_along_axis(double, 1, m)
+        assert isinstance(result, np.matrix)


Note that assert shouldn't be u 10000 sed at all, use np.testing.assert_ instead (applies to multiple lines below as well)

In my defense, those other lines were not part of this patch. Do you want me to fix all cases of assert in this file anyway?

Or should this be self.assertIsInstance?

Changed to assert_(isinstance(...)), because that's what other tests do. Switching to self.assertIsInstance can be discussed elsewhere

In my defense, those other lines were not part of this patch. Do you want me to fix all cases of assert in this file anyway?

if you could, that'd be useful

self.assertIsInstance seems fine to use if you prefer that, it has a sane implementation in both unittest and unittest2.

@rgommers: I've patched the other ones in that file

eric-wieser · 2017-01-20T12:26:12Z

Ugh, this needs a patch in ma.apply_along_axis for consistency as well. Can that wait for a separate PR?

eric-wieser · 2017-01-20T13:07:53Z

numpy/lib/shape_base.py

+    outarr = zeros(outshape, res.dtype)
+    # matrices call reshape in __array_prepare__, and this is apparently ok!
+    if not isinstance(res, matrix):
+        outarr = res.__array_prepare__(outarr)


Can someone comment on whether this is actually correct, or if it can just be omitted altogether

__array_prepare__ seems to always do something screwy for builtin array subclasses. It's not clear to me what the contract entails, and whether I'm violating it, or np.matrix and np.ma.MaskedArray are

@mhvk may know

I think MaskedConstant.__array_prepare__ was broken - fixed in #8508. I think a corollary of this is that np.ma.apply_along_axis can basically defer now.

shoyer

I really like the approach here, thanks for tackling this clean-up!

shoyer · 2017-01-20T17:26:52Z

numpy/lib/shape_base.py

+    outarr = zeros(outshape, res.dtype)
+    # matrices call reshape in __array_prepare__, and this is apparently ok!
+    if not isinstance(res, matrix):
+        outarr = res.__array_prepare__(outarr)


@mhvk may know

shoyer · 2017-01-20T17:30:26Z

numpy/lib/shape_base.py

+    # save the first result
+    outarr_view[ind0] = res
+    for ind in inds:
+        outarr_view[ind] = asanyarray(func1d(inarr_view[ind], *args, **kwargs))


I recognize that this works to update the view inplace, but it still makes me a little nervous. It would be more direct and obviously correct to do the transposing afterwards.

This would change the memory layout of the result array. Transposing afterwards would ensure that array is written with contiguous writes, which could result in a marginal performance improvement (at least for this function).

Yep, my goal here was to leave the memory layout of the result array unaffected when compared to the old version.

In particular, the old code would always give a contiguous output, something that may have been relied upon

I think we reserve the right to change the memory layout of arrays returned by NumPy functions unless it is documented as part of the API. Users who truly rely on contiguous output should be using np.ascontiguousarray.

I guess transposing afterwards would fix the above issue. If you don't think a contiguous output is beneficial, then I'll move the transpose

We have done it before (e.g. even indexing I did it). If it is obvious/possible, then "the same as before" should be preferred most of the time I would say.

I'd argue that the correctness is more obvious this way around, because otherwise the inverse transpose order is needed, which is more troublesome to generate

I mostly like transposing afterwards more. I don't care so much about the memory layout.

Actually I think you're right, it is simpler - the transpose order becomes ~~dims_out[:axis] + dims_out[-res.ndim:] + dims_out[axis+1:]~~ (more complicated than that), which better conveys the insertion of the added axes

Should the transpose happen before or after the __array_wrap__?

Ok, I'm sold on your suggestion, but for a different reason: when res.__array_prepare__ is called, the axes of res (a single output) and the final output are paired as if they were broadcasted, which seems desirable for any kind of axis metadata copying

eric-wieser · 2017-01-20T13:19:56Z

numpy/lib/shape_base.py

+
+    # arr, with the iteration axis at the end
+    dims_in = list(range(nd))
+    inarr_view = transpose(arr, dims_in[:axis] + dims_in[axis+1:] + [axis])


This assumes that transpose always returns a view (and not sometimes a copy). is this acceptable?

(edit: updated mailing list url, which was dead. Not sure it points to the right thing any longer. https://mail.python.org/pipermail/numpy-discussion/2013-June/066822.html might be what I meant)

Actually, only the output transpose does, which only matters if we use __array_prepare__

As @seberg writes in that mailing list discussion, transpose on base numpy arrays always returns a view, never a copy.

But this might indeed break in surprising ways for ndarray subclasses

@shoyer: Would an acceptable compromise be to not call __array_prepare__ if transpose is a copy?

eric-wieser · 2017-01-21T11:29:40Z

@shoyer: Transposed moved, which now gives this masked_array support, leading into #8511

mhvk · 2017-01-21T16:46:33Z

numpy/lib/shape_base.py

    arr = asanyarray(arr)
    nd = arr.ndim
    if axis < 0:
        axis += nd
-    if (axis >= nd):
+    if axis >= nd:


while we're at it, also check that axis is not still negative?

I like the idea of killing this altogether, and delegating parsing axis to moveaxis

Alarmingly, I can't even find a precedent where numpy checks for this anywhere else within shapebase

@eric-wieser - when you add the changelog, could you also move this ch 179B eck up to before axis is possibly changed and check for too-negative as well? That way, the error message gives the actual input rather than something that's already changed. I.e.,

if axis < -nd or axis >= nd: raise ... if axis < 0: axis += nd

@mhvk: Done

mhvk · 2017-01-21T16:56:22Z

numpy/lib/shape_base.py

-        return outarr
+
+    # arr, with the iteration axis at the end
+    in_dims = list(range(nd))


When I looked at this first, I thought you might consider using the (relatively new) np.moveaxis, i.e.,

inarr_view = moveaxis(inarr, axis, -1)

This will not be any faster (as it just sets up a transpose, like here), but is perhaps clearer. As a side benefit, you can remove above the checks on the validity of axis, as moveaxis does those anyway.

However, looking further down, it is not obvious moveaxis would work there, since res.ndim could in principle be >1-d (for which your code adds support, which I think is very nice!). But one could steal a bit from the moveaxis code and write here

in_permute = [n for n in range(nd) if n != axis] + [axis] inarr_view = transpose(arr, in_permute)

Oh, good point. moveaxis works with multiple axes too, so should work in both cases.

For restoring axes, you could maybe write something like:

moveaxis(pre_arr, [i + nd - 1 for i in range(res.ndim)], [i + axis for i in range(res.ndim)])

I think this is a slight improvement in clarity over building up the permutation axes for transpose directly.

I think moveaxis should learn a shorthand for dest such that this would work:

moveaxis(pre_arr, [i + nd - 1 for i in range(res.ndim)], axis)

Ie, if dest is a scalar, move all the source axes to that location, in the order they were passed in src

@shoyer: Based on @mhvk's revised opinion, are you happy for this to remain as it is? Any further pain-points blocking this being merged?

mhvk · 2017-01-21T16:57:13Z

numpy/lib/shape_base.py

+    inarr_view = transpose(arr, in_dims[:axis] + in_dims[axis+1:] + [axis])
+
+    # compute indices for the iteration axes
+    inds = ndindex(inarr_view.shape[:nd-1])


Really like using the iterator here. Much cleaner!

mhvk · 2017-01-21T16:59:15Z

numpy/lib/shape_base.py

+    # remove the requested axis, and add the new ones on the end.
+    # laid out so that each write is contiguous.
+    # for a tuple index inds, pre_arr[inds] = func1d(inarr_view[inds])
+    pre_shape = arr.shape[:axis] + arr.shape[axis+1:] + res.shape


Here you can just use the shape of inarr_view, i.e.,

pre_shape = inarr_view.shape + res.shape

mhvk · 2017-01-21T17:00:33Z

numpy/lib/shape_base.py

+
+    # permutation of axes such that out = pre_arr.transpose(pre_permute)
+    pre_dims = list(range(pre_arr.ndim))
+    pre_permute = pre_dims[:axis] + list(roll(pre_dims[axis:], res.ndim))


Maybe a bit faster (no list->ndarray->list) and clearer (to me at least):

if res.ndim > 0: pre_permute = pre_dims[:axis] + pre_dims[-res.ndim:] + pre_dims[axis+red.ndim:]

(you'd need to rename pre_dims to pre_permute if you want the transpose at the end even for res.ndim=0)

Yeah, that line was inspired by a talk extolling the virtues of std::rotate.

Had that before, but without the if - as you correctly spot, it fails when res.ndim == 0. Not a huge fan of special casing that.

Either way, I reckon moveaxis will do the job here too.

Or how about just

pre_permute = pre_dims[:axis] + pre_dims[nd-res.ndim:] + pre_dims[axis+red.ndim:]

Adding the nd there fixes the error when res.ndim= 0

mhvk · 2017-01-21T17:19:54Z

numpy/lib/shape_base.py

+        pre_arr = res.__array_wrap__(pre_arr)
+
+        # finally, rotate the inserted axes back to where they belong
+        return transpose(pre_arr, pre_permute)


For speed for the more common case (and required with my comment above):

return pre_arr if res.ndim == 0 else transpose(pre_arr, pre_permute)

mhvk · 2017-01-21T17:23:49Z

numpy/lib/shape_base.py

-            outarr = outarr.squeeze(axis)
-        return outarr
+        # matrices have to be transposed first, because they collapse dimensions!
+        out_arr = transpose(pre_arr, pre_permute)


same as above

mhvk · 2017-01-21T17:24:50Z

numpy/lib/tests/test_shape_base.py

        result = apply_along_axis(double, 0, m)
-        assert isinstance(result, np.matrix)
+        assert_(isinstance(result, np.matrix))
+        assert_array_equal(


Just on one line?

assert_array_equal(result, expected)

mhvk · 2017-01-21T17:24:58Z

numpy/lib/tests/test_shape_base.py

        assert_array_equal(
-            result, np.matrix([[0, 2], [4, 6]])
+            result, expected
        )


mhvk · 2017-01-21T17:28:03Z

@eric-wieser - this is very nice indeed, I like the generality, and the care for subclasses. My comments are all small, but hopefully make this a little better still.

eric-wieser · 2017-01-21T19:41:37Z

I'll do another round of clean up tomorrow to address those

mhvk · 2017-01-21T20:11:18Z

I'll do another round of clean up tomorrow to address those

I think I'd go for your approach rather than moveaxis as it is not all that obviously handier in passing in multiple axes.

eric-wieser · 2017-01-22T03:12:54Z

numpy/lib/shape_base.py

@@ -103,11 +103,10 @@ def apply_along_axis(func1d, axis, arr, *args, **kwargs):
            % (axis, nd))

    # arr, with the iteration axis at the end
-    in_dims = list(range(nd))
-    inarr_view = transpose(arr, in_dims[:axis] + in_dims[axis+1:] + [axis])
+    inarr_view = moveaxis(arr, axis, -1)


Beginning to think this would be a fair bit slower, due to the error checking in moveaxis, and isn't much clearer.

Ooh, how about np.r_[:axis, axis+1:nd, axis]

eric-wieser · 2017-01-22T14:43:04Z

Keep changing my mind on this one. I think it's best left without moveaxis. It doesn't add much to clarity in the second case, and I think there's a benefit to having both transposes written in the same style. Also, building the axis lists directly should be a lot faster than incurring list comprehensions, then the type conversion and error checking that occurs in moveaxis.

np.r_ added some minor conciseness, but at the expense of some performance

I agree that the roll was a little cryptic in the second case.

I'd rather not spend too much more time bikeshedding this.

mhvk · 2017-01-22T15:51:33Z

@eric-wieser - sorry for having brought up moveaxis when I had also concluded that in the end it was not as good an idea -- it works well for the first case, not so well for the second.

shoyer · 2017-02-08T17:39:51Z

numpy/lib/tests/test_shape_base.py

+
+    def test_axis_insertion(self, cls=np.ndarray):
+        a = np.arange(18).reshape((6, 3))
+        res = apply_along_axis(lambda x: np.diag(x).view(cls), 0, a)


Reading over these tests one last time, it occurs to me that we should be more comprehensive for testing insertion of new axes.

Right now you only test inserting at the start, but that doesn't really exercise the permutation logic in a comprehensive way. Let's also test inserting in the middle and at the end. Also, diag is axes order invariant, so none of the tests verify that new axes are inserted in the right order.

Maybe outer product with its reverse would be better than diag then?

I think the outer product is still invariant to transposes. Something like x[:, np.newaxis] - x[np.newaxis, :] would work, though.

On re-reading: yes, outer product with its reverse, would be fine

You didn't miss it the first time, I did indeed suggest a plain outer product beforehand, and realized my mistake. This is now done too.

shoyer · 2017-02-08T17:41:31Z

numpy/lib/tests/test_shape_base.py

+        res = apply_along_axis(lambda x: np.diag(x).view(cls), 0, a)
+        assert_(isinstance(res, cls))
+        assert_equal(res.ndim, 3)
+        assert_array_equal(res[:,:,0], np.diag(a[:,0]).view(cls))


A cleaner way to check this sort of thing is to construct the desired result array and just call assert_array_equal (e.g., np.stack([np.diag(a[:,i]) for i in range(3)], axis=-1).view(cls)). Otherwise it's easy to overlook one part of the equality check (e.g., shape, in this case).

@shoyer: I was hoping I could be done with this, but you are definitely right. I'll fix those things up soon

Ok, done. (except for masked arrays, where I don't think assert_equal does the right thing)

shoyer

LGTM. Will merge this shortly if no one else has further comments.

shoyer · 2017-02-10T18:53:17Z

@eric-wieser could you please clean up the git history a little bit? Thanks

eric-wieser · 2017-02-11T01:08:18Z

@shoyer: Fair, I'll give that a go over the weekend.

Also: ENH: Support arbitrary dimensionality of return value MAINT: remove special casing

.transpose does not specify that it must return a view, so subclasses (like np.ma.array) could otherwise break this. This exposes some more need for matrix special casing.

Note that this is not a full subsitute for np.ma.apply_along_axis, as that allows functions to return a mix of np.ma.masked and scalars

Copied from the implementation in core.shape_base.stack

eric-wieser · 2017-02-11T21:14:39Z

@shoyer: Down from 13 commits to 8 commits. Look ok?

I've confirmed that the final commit of these squashes is identical to the result of rebasing this branch on master. (possibly ignoring merges in the release notes)

shoyer · 2017-02-12T07:50:39Z

I would usually squash down to one commit, but this is fine.

eric-wieser force-pushed the apply_along_axis-nd branch from 3c6ea94 to daa4f4e Compare January 2, 2017 23:18

eric-wieser commented Jan 2, 2017

View reviewed changes

eric-wieser force-pushed the apply_along_axis-nd branch from f71bedf to 27cd75b Compare January 2, 2017 23:39

eric-wieser commented Jan 2, 2017

View reviewed changes

eric-wieser force-pushed the apply_along_axis-nd branch from 8d761e9 to 3a159d4 Compare January 3, 2017 01:29

eric-wieser mentioned this pull request Jan 3, 2017

BUG: Get common dtype in apply_along_axis #8363

Closed

eric-wieser commented Jan 3, 2017

View reviewed changes

rgommers reviewed Jan 3, 2017

View reviewed changes

rgommers added 00 - Bug component: numpy.lib labels Jan 7, 2017

eric-wieser mentioned this pull request Jan 20, 2017

__array_prepare__ produces bad shape for np.ma.masked #8505

Closed

eric-wieser commented Jan 20, 2017

View reviewed changes

shoyer reviewed Jan 20, 2017

View reviewed changes

eric-wieser commented Jan 20, 2017

View reviewed changes

eric-wieser force-pushed the apply_along_axis-nd branch from ce5a18e to 8dba2ba Compare January 20, 2017 22:42

eric-wieser mentioned this pull request Jan 21, 2017

MAINT: make np.ma.apply_along_axis consistent with np.apply_along_axis #8511

Closed

mhvk reviewed Jan 21, 2017

View reviewed changes

eric-wieser commented Jan 22, 2017

View reviewed changes

eric-wieser force-pushed the apply_along_axis-nd branch 2 times, most recently from 060c5e7 to 7939511 Compare January 22, 2017 03:35

eric-wieser force-pushed the apply_along_axis-nd branch from 7939511 to 7f4ab47 Compare January 22, 2017 15:08

shoyer reviewed Feb 8, 2017

View reviewed changes

shoyer approved these changes Feb 9, 2017

View reviewed changes

eric-wieser added 8 commits February 11, 2017 21:08

BUG: Fix crash on 0d return value in apply_along_axis

52988ea

Also: ENH: Support arbitrary dimensionality of return value MAINT: remove special casing

MAINT: Use np.ndindex, which seems just as efficient

5307aed

BUG: Call __array_prepare__ before __array_wrap__

b10b6c2

BUG: Work around evil matrix.__array_prepare__

9f362bf

MAINT: Transpose the result, rather than working with a transposed view

c8efc57

.transpose does not specify that it must return a view, so subclasses (like np.ma.array) could otherwise break this. This exposes some more need for matrix special casing.

TST: Verify apply_along_axis now works on masked arrays

ff9c363

Note that this is not a full subsitute for np.ma.apply_along_axis, as that allows functions to return a mix of np.ma.masked and scalars

MAINT: Improve error-checking of axis argument

78084ee

Copied from the implementation in core.shape_base.stack

DOC: Update 1.13.0 release notes with apply_along_axis changes

d4bce01

eric-wieser force-pushed the apply_along_axis-nd branch from 5435697 to d4bce01 Compare February 11, 2017 21:13

shoyer merged commit 51a8240 into numpy:master Feb 12, 2017

homu mentioned this pull request Feb 12, 2017

MAINT: Use the same exception for all bad axis requests #8584

Merged

lesteve mentioned this pull request Feb 14, 2017

[MRG] Fix tests on numpy master scikit-learn/scikit-learn#8355

Merged

eric-wieser mentioned this pull request Mar 1, 2017

ENH: Implement take_along_axis as described in #8708 #8714

Closed

eric-wieser mentioned this pull request Sep 5, 2018

DOC: Recommend the use of np.ndim over np.isscalar, and explain the differences #11882

Merged

eric-wieser mentioned this pull request Dec 3, 2019

DOC: clarify documentation for transpose() #15024

Merged

Uh oh!

BUG: Fix crash on 0d return value in apply_along_axis #8441

BUG: Fix crash on 0d return value in apply_along_axis #8441

Uh oh!

Conversation

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!