Repository: yt-project/unyt

Issue:
BUG: `linspace` with `retstep` fails
* unyt version: 3.0.3
* Python version: 3.12
* Operating System: Linux (Fedora 40)

### Description

The `linspace` function is not implemented to expect and handle multiple return values when the `retstep` keyword argument is set to `True`.

### What I Did

```
>>> np.linspace(0 * u.cm, 1 * u.cm, num=5, retstep=True)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/txwx36/code/unyt/unyt/array.py", line 2045, in __array_function__
    return _HANDLED_FUNCTIONS[func](*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/txwx36/code/unyt/unyt/_array_functions.py", line 497, in linspace
    np.linspace._implementation(
  File "/home/txwx36/code/unyt/unyt/unit_object.py", line 374, in __rmul__
    return self.__mul__(u)
           ^^^^^^^^^^^^^^^
  File "/home/txwx36/code/unyt/unyt/unit_object.py", line 379, in __mul__
    data = np.array(u, subok=True)
           ^^^^^^^^^^^^^^^^^^^^^^^
ValueError: setting an array element with a sequence. The requested array has an inhomogeneous shape after 1 dimensions. The detected shape was (2,) + inhomogeneous part.
```

Discussion:

Solution patch:
diff --git a/unyt/_array_functions.py b/unyt/_array_functions.py
index 3bd811b7..c8e92f97 100644
--- a/unyt/_array_functions.py
+++ b/unyt/_array_functions.py
@@ -621,26 +621,94 @@ def array_equiv(a1, a2, *args, **kwargs) -> bool:
     )
 
 
-@implements(np.linspace)
-def linspace(start, stop, *args, **kwargs):
+def _linspace(
+    start,
+    stop,
+    num=50,
+    endpoint=True,
+    retstep=False,
+    dtype=None,
+    axis=0,
+    *,
+    device=None,
+):
     _validate_units_consistency((start, stop))
-    return (
-        np.linspace._implementation(
-            np.asarray(start), np.asarray(stop), *args, **kwargs
+    kwargs = {
+        "num": num,
+        "endpoint": endpoint,
+        "retstep": retstep,
+        "dtype": dtype,
+        "axis": axis,
+    }
+    if NUMPY_VERSION >= Version("2.1.0.dev0"):
+        kwargs["device"] = device
+    result = np.linspace._implementation(np.asarray(start), np.asarray(stop), **kwargs)
+    if retstep:
+        return result[0] * start.units, result[1] * start.units
+    else:
+        return result * start.units
+
+
+if NUMPY_VERSION >= Version("2.0.0.dev0"):
+
+    @implements(np.linspace)
+    def linspace(
+        start,
+        stop,
+        num=50,
+        endpoint=True,
+        retstep=False,
+        dtype=None,
+        axis=0,
+        *,
+        device=None,
+    ):
+        return _linspace(
+            start,
+            stop,
+            num=50,
+            endpoint=endpoint,
+            retstep=retstep,
+            dtype=dtype,
+            axis=axis,
+            device=device,
+        )
+
+else:
+
+    @implements(np.linspace)
+    def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None, axis=0):
+        return _linspace(
+            start,
+            stop,
+            num=50,
+            endpoint=endpoint,
+            retstep=retstep,
+            dtype=dtype,
+            axis=axis,
         )
-        * start.units
-    )
 
 
 @implements(np.logspace)
-def logspace(start, stop, *args, **kwargs):
-    _validate_units_consistency((start, stop))
-    return (
-        np.logspace._implementation(
-            np.asarray(start), np.asarray(stop), *args, **kwargs
+def logspace(start, stop, num=50, endpoint=True, base=10, dtype=None, axis=0):
+    startu = getattr(start, "units", NULL_UNIT)
+    stopu = getattr(stop, "units", NULL_UNIT)
+    if startu != NULL_UNIT or stopu != NULL_UNIT:
+        raise TypeError(
+            "The first two arguments to numpy.logspace must be dimensionless, "
+            f"got units={startu} (arg1) and units={stopu} (arg2). If output with"
+            " units is desired, apply the units to the `base` kwarg."
         )
-        * start.units
+    result = np.logspace._implementation(
+        np.asarray(start),
+        np.asarray(stop),
+        num=num,
+        endpoint=endpoint,
+        base=np.asarray(base),
+        dtype=dtype,
+        axis=axis,
     )
+    return result * getattr(base, "units", NULL_UNIT)
 
 
 @implements(np.geomspace)


Test patch:
diff --git a/unyt/tests/test_array_functions.py b/unyt/tests/test_array_functions.py
index 291db7d8..fbabac1d 100644
--- a/unyt/tests/test_array_functions.py
+++ b/unyt/tests/test_array_functions.py
@@ -9,7 +9,7 @@
 from numpy.testing import assert_allclose
 from packaging.version import Version
 
-from unyt import A, K, Msun, cm, degC, delta_degC, g, km, rad, s
+from unyt import A, K, Msun, cm, degC, delta_degC, dimensionless, g, km, rad, s
 from unyt._array_functions import (
     _HANDLED_FUNCTIONS as HANDLED_FUNCTIONS,
     _UNSUPPORTED_FUNCTIONS as UNSUPPORTED_FUNCTIONS,
@@ -1233,16 +1233,42 @@ def test_isclose_error():
         np.isclose(x, y)
 
 
-@pytest.mark.parametrize(
-    "func",
-    [
-        np.linspace,
-        np.logspace,
-        np.geomspace,
-    ],
-)
-def test_xspace(func):
-    res = func(1 * cm, 11 * cm, 10)
+def test_linspace():
+    res = np.linspace(1 * cm, 11 * cm, 10)
+    assert type(res) is unyt_array
+    assert res.units == cm
+
+
+def test_linspace_with_retstep():
+    res, step = np.linspace(1 * cm, 11 * cm, 10, retstep=True)
+    assert type(res) is unyt_array
+    assert res.units == cm
+    assert type(step) is unyt_quantity
+    assert step.units == cm
+
+
+def test_logspace_with_units_raises():
+    with pytest.raises(
+        TypeError,
+        match="The first two arguments to numpy.logspace must be dimensionless",
+    ):
+        np.logspace(1 * cm, 2 * cm)
+
+
+def test_logspace_with_base():
+    res = np.logspace(1, 2, base=10 * cm)
+    assert type(res) is unyt_array
+    assert res.units == cm
+
+
+def test_logspace_with_dimless_input():
+    res = np.logspace(1 * dimensionless, 2 * dimensionless)
+    assert type(res) is unyt_array
+    assert res.units == dimensionless
+
+
+def test_geomspace():
+    res = np.geomspace(1 * cm, 11 * cm, 10)
     assert type(res) is unyt_array
     assert res.units == cm