First sentence should fit on a single line. Maybe

Take elements from slices indexed along the given axis.

The problem is I need to disambiguate this from take, which is

Take elements from an array along an axis.

Should go to the mailing list in any case, maybe take is not even the most natural word in the end, which would somewhat remove the problem ;). The take description kind of only works for 1-D, and in 1-D the two do the same thing, so its a bit of a twist :).

It doesn't have to be my example, but I do think it should be a single line.

You're not wrong - I'm just asking for help in coming up with an unambiguous description under that constraint :)

The only other idea I have right now is to call it a vectorized take/pick (which bites with the vindex idea, but maybe that name is not great anyway, I think someone had suggest broadcasted index there too, which may be more logical anyway -- though that does not necessarily mean easier I guess, hehe).

Unify use of notation for indices (e.g., here lowercase a, b; below upper case; ideally use the standard "integers", i.e., i--n). If possible, do use axis directly.

Here I'm using the convention that a = range(0, A), ie the capital letters are the shape, and the lowercase ones indices for that dimension

There's no way I can use axis directly here, short of some ascii art pointing to the middle index on the right

I could use out[i..., j..., k...] = arr[i..., indices[i..., j..., k...], k...], and then Ni, Nk, Nk further down?

I like the i,j,k, Ni,Nj,Nk, or perhaps i1, i2, i3, N1,N2,N3.

I've added a fixup commit to apply this. I'll squash once everything else is approved

Could you add an example where one keeps the dimension?

Again, I think keeping dimensions is out of scope here, and belongs in #8710.

"keeping the dimension" is something that can be done either before or after take_along_axis.

In principle we could think about making keepdims (well, kind of the inverse) a kwarg here too.

So you are lazy about a C version, too bad ;P

@seberg: What would it do though? Dimensions are already kept, in that out.ndim == indices.ndim.

But there is no problem with argmax needing an expand dims, since we should just add a keepdims?

Right - it's inconvenient right now to have to call np.expand_dims, but that's really a bug of not having keepdims in argmethods. And as for calling squeeze on the return value (vs what I proposed before) - you're likely in for a bad time if you start squeezing axes in the middle of your multidimensional data anyway.

Having assert(ind.ndim == x.ndim) still gives us plenty of freedom to come up with better semantics in future

Another example of how this restriction is sufficient (with keepdims) - indexing the brightest pixel:

assert img.shape == (400, 300, 3)
brightness = np.sum(img, axis=2, keepdims=True)  # again - always keep your dims!
argbrightrow = np.argmax(brightness, axis=1, keepdims=True)
brightest_by_row = np.take_along_axis(img, argbrightrow, axis=1)

Having assert(ind.ndim == x.ndim) still gives us plenty of freedom to come up with better semantics in future

We could remove this and add the default broadcasting behaviour, but this risks confusing users who try to do np.take_along_axis(a, a.argmax(axis=1),axis=1),(no keepdim) which would allow this to sometimes silently do the wrong thing

So for now, I guess lets just say we don't put it in, and put it up for discussion on the mailing list. If anyone can present a good reason/usecase especially for the way you first described it, I am good with allowing it. At least the 0-D case (dim is missing), it seems rather intuitive after all....

I would not use ; in example code; just make another line...

This is pretty par for the course for numpy docstrings - grep for >>> (\w+)\s*=.*;\s*\1

Bad practice anyway.

Resolved

This should be a TypeError, I think (at least, [1, 2][1.] raises TypeError).

Just copying np.array([1, 2])[np.array(1.)] here, which gives IndexError. This is just that error message, but without the bit about booleans.

Bools are not considered ints here right?

Correct, np.issubdtype(np.bool_, np.integer) is false. There's a test for this error in this PR.

OK, should have tried with arrays; not very logical but best to stick with numpy practice here.

Just to be sure, also add the keepdims versions.

I'd argue for postponing that to #8710 as well, and then no special handling would be needed.
I have added a test to verify that expand_dims works before and after though, which is sort of the same thing

Based on the discussion #9946, I suppose this would be better described as

Ni, Nk = a.shape[:axis], a.shape[axis+1:]
for ii in ndindex(Ni):
    for kk in ndindex(Nj):
        out[ii + s_[...,] + kk] = a[ii + s_[:,] + kk][indices[ii + s_[...,] + kk]]

Edit: updated

Could you add support for arr.ndim > indices.ndim?
It's a single line here (+ unit test + docs):

indices = indices.reshape((1, ) * (arr.ndim - indices.ndim) + indices.shape)

Use case / example:

x = np.array([[5, 3, 2, 8, 1],
              [0, 7, 1, 3, 2]])
# Completely arbitrary y = f(x0, x1, ..., xn), embarassingly parallel along axis=-1
# Here we only have x0, but we could have more.
y = x.sum(axis=0)
# Sort the x's, moving the ones that cause the smallest y's to the left
take_along_axis(x, np.argsort(y))

I don't think your use-case is well-motivated. A more explicit way to achieve that would be:

y = x.sum(axis=0, keepdims=True)
take_along_axis(x, np.argsort(y, axis=1), axis=1)

The whole point is that with the one-liner addition f(x) can add or remove axes at will (some manual broadcasting required if it replaces or transposes axes, which however happens automatically if you e.g. wrap this in xarray.apply_ufunc).

Note that in the comments above, we decide that perhaps it's best to not allow any case other than indices.ndim == arr.ndim, since there's no obvious right choice.

take_along_axis only really makes sense if you endeavor to keep all your axes aligned. xarray can probably solve that by axis names alone, but in numpy you need to indicate that by axis position. Therefore, you can't afford to let your axes collapse, and have numpy guess which one you lost: in your case, you're advocating for it to guess the left-most one should be reinserted - but this is only the case because you did sum(axis=0).

default to axis=-1 like 99% of other numpy functions?

default to axis=-1 like 99% of other numpy functions?

I wish that were true. concatenate defaults to axis=0, and other functions default to axis=None. Since the axis is a key part of the function, it seems best just to require it.