Repository: yt-project/unyt

Issue:
BUG: numpy function `take` strips units with scalar index input
* 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)
```


Discussion:
kyleaoman: pre-commit.ci autofix
kyleaoman: pre-commit.ci autofix

Patch:
diff --git a/unyt/_array_functions.py b/unyt/_array_functions.py
index 03a4875..b32d15f 100644
--- a/unyt/_array_functions.py
+++ b/unyt/_array_functions.py
@@ -1220,3 +1220,23 @@ if hasattr(np, "in1d"):
         return np.isin._implementation(
             np.asarray(ar1), np.asarray(ar2), *args, **kwargs
         )
+
+
+@implements(np.take)
+def take(a, indices, axis=None, out=None, mode="raise"):
+    ret_units = getattr(a, "units", NULL_UNIT)
+
+    if out is not None:
+        out_view = np.asarray(out)
+    else:
+        out_view = None
+
+    res = np.take._implementation(
+        np.asarray(a), indices, axis=axis, out=out_view, mode=mode
+    )
+
+    if getattr(out, "units", None) is not None:
+        out.units = ret_units
+
+    ret_cls = unyt_quantity if res.ndim == 0 else unyt_array
+    return ret_cls(res, ret_units, bypass_validation=True)
diff --git a/unyt/array.py b/unyt/array.py
index b2699ef..ff9dbde 100644
--- a/unyt/array.py
+++ b/unyt/array.py
@@ -2125,6 +2125,21 @@ class unyt_array(np.ndarray):
             out.units = res_units
         return ret
 
+    def take(self, indices, axis=None, out=None, mode="raise"):
+        """method
+
+        Return an array formed from the elements of `a` at the given indices.
+
+        Refer to :func:`numpy.take` for full documentation.
+
+        See also
+        --------
+        numpy.take : equivalent function
+        """
+        from ._array_functions import take
+
+        return take(self, indices, axis=axis, out=out, mode=mode)
+
     def __reduce__(self):
         """Pickle reduction method
 


Test patch:
diff --git a/unyt/tests/test_array_functions.py b/unyt/tests/test_array_functions.py
index a8b582c..3da71ef 100644
--- a/unyt/tests/test_array_functions.py
+++ b/unyt/tests/test_array_functions.py
@@ -143,7 +143,6 @@ NOOP_FUNCTIONS = {
     np.nancumprod,  # we get it for free with np.cumprod (tested)
     np.bincount,  # works out of the box (tested)
     np.unique,  # works out of the box (tested)
-    np.take,  # works out of the box (tested)
     np.min_scalar_type,  # returns dtypes
     np.extract,  # works out of the box (tested)
     np.setxor1d,  # we get it for free with previously implemented functions (tested)
@@ -1848,10 +1847,25 @@ def test_unique_values():
     assert_array_equal_units(res, values)
 
 
-def test_take():
+@pytest.mark.parametrize("indices", [[0, 1], 0])
+def test_take(indices):
     a = [1, 2, 3] * cm
-    res = np.take(a, [0, 1])
-    assert type(res) is unyt_array
+    res = np.take(a, indices)
+    if res.ndim == 0:
+        assert type(res) is unyt_quantity
+    else:
+        assert type(res) is unyt_array
+    assert res.units == cm
+
+
+@pytest.mark.parametrize("indices", [[0, 1], 0])
+def test_ndarray_take(indices):
+    a = [1, 2, 3] * cm
+    res = a.take(indices)
+    if res.ndim == 0:
+        assert type(res) is unyt_quantity
+    else:
+        assert type(res) is unyt_array
     assert res.units == cm