BUG: numpy function take strips units with scalar index input

[kyleaoman]

• unyt version: 3.0.3
• Python version: 3.12
• Operating System: Linux (Fedora 40)

Description

The numpy function take strips units when a scalar index (rather than array-like) is passed.

What I Did

>>> a = [1, 2, 3] * u.cm
>>> a.take([0])  # OK
unyt_array([1], 'cm')
>>> np.take(a, [0])  # OK
unyt_array([1], 'cm')
>>> a.take(0)  # not OK
np.int64(1)
>>> np.take(a, 0)  # not OK
np.int64(1)

[neutrinoceros]

Ah, this is a case where I just wrote the simplest test case possible (test_take), saw that it worked, and went to the next function. So, we actually don't have an implementation for np.take at the moment.
A quick test reveals that implementing np.take won't automatically make unyt_array.take correct either, so I think we also need to override np.ndarray.take here. We could probably implement one using the other though. Do you want to give it a go ?

[kyleaoman]

Yes possibly, but I'm now on xmas holiday, so maybe not soon :)

[neutrinoceros]

No worries. Happy holidays !