Repository: yt-project/unyt
Pull request #553 (2024-12-25 17:19:53)

BUG: fix a regression where `np.linspace`'s num argument would be ignored for `unyt_array`s
This fixes a regression I overlooked when reviewing #544 but poped up in yt's test suite https://github.com/yt-project/yt/issues/5091

for posterity: this regression never landed in a release.


Patch:
diff --git a/unyt/_array_functions.py b/unyt/_array_functions.py
index b32d15f..df14564 100644
--- a/unyt/_array_functions.py
+++ b/unyt/_array_functions.py
@@ -666,7 +666,7 @@ if NUMPY_VERSION >= Version("2.0.0.dev0"):
         return _linspace(
             start,
             stop,
-            num=50,
+            num=num,
             endpoint=endpoint,
             retstep=retstep,
             dtype=dtype,
@@ -681,7 +681,7 @@ else:
         return _linspace(
             start,
             stop,
-            num=50,
+            num=num,
             endpoint=endpoint,
             retstep=retstep,
             dtype=dtype,
@@ -1220,23 +1220,3 @@ 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 ff9dbde..b2699ef 100644
--- a/unyt/array.py
+++ b/unyt/array.py
@@ -2125,21 +2125,6 @@ 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 3da71ef..a8b582c 100644
--- a/unyt/tests/test_array_functions.py
+++ b/unyt/tests/test_array_functions.py
@@ -143,6 +143,7 @@ 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)
@@ -1847,25 +1848,10 @@ def test_unique_values():
     assert_array_equal_units(res, values)
 
 
-@pytest.mark.parametrize("indices", [[0, 1], 0])
-def test_take(indices):
+def test_take():
     a = [1, 2, 3] * cm
-    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
+    res = np.take(a, [0, 1])
+    assert type(res) is unyt_array
     assert res.units == cm