Improve linspace implementation with non-scalar inputs (#2712)

The PR closes #2084.

This PR improves implementation of `dpnp.linspace` function aligning it
with the latest changes added to NumPy.
The tests coverage is also extended to verify more use cases.

Author: antonwolfy

CHANGELOG.md

@@ -59,6 +59,7 @@ Also, that release drops support for Python 3.9, making Python 3.10 the minimum
 
 * Suppressed a potential deprecation warning triggered during import of the `dpctl.tensor` module [#2709](https://github.com/IntelPython/dpnp/pull/2709)
 * Corrected a phonetic spelling issue due to incorrect using of `a nd` in docstrings [#2719](https://github.com/IntelPython/dpnp/pull/2719)
+* Resolved an issue causing `dpnp.linspace` to return an incorrect output shape when inputs were passed as arrays [#2712](https://github.com/IntelPython/dpnp/pull/2712)
 
 ### Security
 

dpnp/dpnp_algo/dpnp_arraycreation.py

@@ -172,14 +172,9 @@ def dpnp_linspace(
 
     num = operator.index(num)
     if num < 0:
-        raise ValueError("Number of points must be non-negative")
+        raise ValueError(f"Number of samples={num} must be non-negative.")
     step_num = (num - 1) if endpoint else num
 
-    step_nan = False
-    if step_num == 0:
-        step_nan = True
-        step = dpnp.nan
-
     if dpnp.isscalar(start) and dpnp.isscalar(stop):
         # Call linspace() function for scalars.
         usm_res = dpt.linspace(
@@ -191,8 +186,13 @@ def dpnp_linspace(
             sycl_queue=sycl_queue_normalized,
             endpoint=endpoint,
         )
-        if retstep is True and step_nan is False:
-            step = (stop - start) / step_num
+
+        # calculate the used step to return
+        if retstep is True:
+            if step_num > 0:
+                step = (stop - start) / step_num
+            else:
+                step = dpnp.nan
     else:
         usm_start = dpt.asarray(
             start,
@@ -204,6 +204,8 @@ def dpnp_linspace(
             stop, dtype=dt, usm_type=_usm_type, sycl_queue=sycl_queue_normalized
         )
 
+        delta = usm_stop - usm_start
+
         usm_res = dpt.arange(
             0,
             stop=num,
@@ -212,20 +214,30 @@ def dpnp_linspace(
             usm_type=_usm_type,
             sycl_queue=sycl_queue_normalized,
         )
+        usm_res = dpt.reshape(usm_res, (-1,) + (1,) * delta.ndim, copy=False)
+
+        if step_num > 0:
+            step = delta / step_num
+
+            # Needed a special handling for denormal numbers (when step == 0),
+            # see numpy#5437 for more details.
+            # Note, dpt.where() is used to avoid a synchronization branch.
+            usm_res = dpt.where(
+                step == 0, (usm_res / step_num) * delta, usm_res * step
+            )
+        else:
+            step = dpnp.nan
+            usm_res = usm_res * delta
 
-        if step_nan is False:
-            step = (usm_stop - usm_start) / step_num
-            usm_res = dpt.reshape(usm_res, (-1,) + (1,) * step.ndim, copy=False)
-            usm_res = usm_res * step
-            usm_res += usm_start
+        usm_res += usm_start
 
         if endpoint and num > 1:
-            usm_res[-1] = dpt.full(step.shape, usm_stop)
+            usm_res[-1, ...] = usm_stop
 
     if axis != 0:
         usm_res = dpt.moveaxis(usm_res, 0, axis)
 
-    if numpy.issubdtype(dtype, dpnp.integer):
+    if dpnp.issubdtype(dtype, dpnp.integer):
         dpt.floor(usm_res, out=usm_res)
 
     res = dpt.astype(usm_res, dtype, copy=False)

dpnp/dpnp_iface_arraycreation.py

@@ -2704,6 +2704,8 @@ def linspace(
         of tuples, tuples of lists, and ndarrays. If `endpoint` is set to
         ``False`` the sequence consists of all but the last of ``num + 1``
         evenly spaced samples, so that `stop` is excluded.
+    num : int
+        Number of samples. Must have a nonnegative value.
     dtype : {None, str, dtype object}, optional
         The desired dtype for the array. If not given, a default dtype will be
         used that can represent the values (by considering Promotion Type Rule

dpnp/tests/helper.py

@@ -69,6 +69,9 @@ def assert_dtype_allclose(
         x.dtype, dpnp.inexact
     )
 
+    if not hasattr(numpy_arr, "dtype"):
+        numpy_arr = numpy.array(numpy_arr)
+
     if is_inexact(dpnp_arr) or is_inexact(numpy_arr):
         tol_dpnp = (
             dpnp.finfo(dpnp_arr).resolution

dpnp/tests/test_arraycreation.py

@@ -19,6 +19,8 @@
     assert_dtype_allclose,
     get_all_dtypes,
     get_array,
+    get_float_dtypes,
+    has_support_aspect64,
     is_lts_driver,
     is_tgllp_iris_xe,
     is_win_platform,
@@ -83,6 +85,131 @@ def test_validate_positional_args(self, xp):
         )
 
 
+class TestLinspace:
+    @pytest.mark.parametrize("start", [0, -5, 10, -2.5, 9.7])
+    @pytest.mark.parametrize("stop", [0, 10, -2, 20.5, 120])
+    @pytest.mark.parametrize("num", [0, 1, 5, numpy.array(10)])
+    @pytest.mark.parametrize(
+        "dt", get_all_dtypes(no_bool=True, no_float16=False)
+    )
+    @pytest.mark.parametrize("retstep", [True, False])
+    def test_basic(self, start, stop, num, dt, retstep):
+        if (
+            not has_support_aspect64()
+            and numpy.issubdtype(dt, numpy.integer)
+            and start == -5
+            and stop == 10
+            and num == 10
+        ):
+            pytest.skip("due to dpctl-1056")
+
+        if numpy.issubdtype(dt, numpy.unsignedinteger):
+            start = abs(start)
+            stop = abs(stop)
+
+        res = dpnp.linspace(start, stop, num, dtype=dt, retstep=retstep)
+        exp = numpy.linspace(start, stop, num, dtype=dt, retstep=retstep)
+        if retstep:
+            res, res_step = res
+            exp, exp_step = exp
+            assert_dtype_allclose(res_step, exp_step)
+
+        if numpy.issubdtype(dt, numpy.integer):
+            assert_allclose(res, exp, rtol=1)
+        else:
+            assert_dtype_allclose(res, exp)
+
+    @pytest.mark.parametrize(
+        "start, stop",
+        [
+            (dpnp.array(1), dpnp.array([-4])),
+            (dpnp.array([2.6]), dpnp.array([[2.6], [-4]])),
+            (numpy.array([[-6.7, 3]]), numpy.array(2)),
+            ([1, -4], [[-4.6]]),
+            ((3, 5), (3,)),
+        ],
+    )
+    @pytest.mark.parametrize("num", [0, 1, 5])
+    @pytest.mark.parametrize(
+        "dt", get_all_dtypes(no_bool=True, no_float16=False)
+    )
+    @pytest.mark.parametrize("retstep", [True, False])
+    def test_start_stop_arrays(self, start, stop, num, dt, retstep):
+        res = dpnp.linspace(start, stop, num, dtype=dt, retstep=retstep)
+        exp = numpy.linspace(
+            get_array(numpy, start),
+            get_array(numpy, stop),
+            num,
+            dtype=dt,
+            retstep=retstep,
+        )
+        if retstep:
+            res, res_step = res
+            exp, exp_step = exp
+            assert_dtype_allclose(res_step, exp_step)
+        assert_dtype_allclose(res, exp)
+
+    @pytest.mark.parametrize(
+        "start, stop",
+        [(1 + 2j, 3 + 4j), (1j, 10), ([0, 1], 3 + 2j)],
+    )
+    def test_start_stop_complex(self, start, stop):
+        result = dpnp.linspace(start, stop, num=5)
+        expected = numpy.linspace(start, stop, num=5)
+        assert_dtype_allclose(result, expected)
+
+    @pytest.mark.parametrize("dt", get_float_dtypes())
+    def test_denormal_numbers(self, dt):
+        stop = numpy.nextafter(dt(0), dt(1)) * 5  # denormal number
+
+        result = dpnp.linspace(0, stop, num=10, endpoint=False, dtype=dt)
+        expected = numpy.linspace(0, stop, num=10, endpoint=False, dtype=dt)
+        assert_dtype_allclose(result, expected)
+
+    @pytest.mark.skipif(not has_support_aspect64(), reason="due to dpctl-1056")
+    def test_equivalent_to_arange(self):
+        result = dpnp.linspace(0, 35, num=36, dtype=int)
+        expected = numpy.linspace(0, 35, num=36, dtype=int)
+        assert_equal(result, expected)
+
+    def test_round_negative(self):
+        result = dpnp.linspace(-1, 3, num=8, dtype=int)
+        expected = numpy.linspace(-1, 3, num=8, dtype=int)
+        assert_array_equal(result, expected)
+
+    def test_step_zero(self):
+        start = numpy.array([0.0, 1.0])
+        stop = numpy.array([2.0, 1.0])
+
+        result = dpnp.linspace(start, stop, num=3)
+        expected = numpy.linspace(start, stop, num=3)
+        assert_array_equal(result, expected)
+
+    @pytest.mark.parametrize("endpoint", [True, False])
+    def test_num_zero(self, endpoint):
+        start, stop = 0, [0, 1, 2, 3, 4]
+        result = dpnp.linspace(start, stop, num=0, endpoint=endpoint)
+        expected = numpy.linspace(start, stop, num=0, endpoint=endpoint)
+        assert_dtype_allclose(result, expected)
+
+    @pytest.mark.parametrize("axis", [0, 1])
+    def test_axis(self, axis):
+        func = lambda xp: xp.linspace([2, 3], [20, 15], num=10, axis=axis)
+        assert_allclose(func(dpnp), func(numpy))
+
+    @pytest.mark.parametrize("xp", [dpnp, numpy])
+    def test_negative_num(self, xp):
+        with pytest.raises(ValueError, match="must be non-negative"):
+            _ = xp.linspace(0, 10, num=-1)
+
+    @pytest.mark.parametrize("xp", [dpnp, numpy])
+    def test_float_num(self, xp):
+        with pytest.raises(
+            TypeError, match="cannot be interpreted as an integer"
+        ):
+            _ = xp.linspace(0, 1, num=2.5)
+
+
 class TestTrace:
     @pytest.mark.parametrize("a_sh", [(3, 4), (2, 2, 2)])
     @pytest.mark.parametrize(
@@ -734,37 +861,6 @@ def test_dpctl_tensor_input(func, args):
         assert_array_equal(X, Y)
 
 
-@pytest.mark.parametrize("start", [0, -5, 10, -2.5, 9.7])
-@pytest.mark.parametrize("stop", [0, 10, -2, 20.5, 120])
-@pytest.mark.parametrize(
-    "num",
-    [1, 5, numpy.array(10), dpnp.array(17), dpt.asarray(100)],
-    ids=["1", "5", "numpy.array(10)", "dpnp.array(17)", "dpt.asarray(100)"],
-)
-@pytest.mark.parametrize(
-    "dtype",
-    get_all_dtypes(no_bool=True, no_float16=False),
-)
-@pytest.mark.parametrize("retstep", [True, False])
-def test_linspace(start, stop, num, dtype, retstep):
-    if numpy.issubdtype(dtype, numpy.unsignedinteger):
-        start = abs(start)
-        stop = abs(stop)
-
-    res_np = numpy.linspace(start, stop, num, dtype=dtype, retstep=retstep)
-    res_dp = dpnp.linspace(start, stop, num, dtype=dtype, retstep=retstep)
-
-    if retstep:
-        [res_np, step_np] = res_np
-        [res_dp, step_dp] = res_dp
-        assert_allclose(step_np, step_dp)
-
-    if numpy.issubdtype(dtype, dpnp.integer):
-        assert_allclose(res_np, res_dp, rtol=1)
-    else:
-        assert_dtype_allclose(res_dp, res_np)
-
-
 @pytest.mark.parametrize("func", ["geomspace", "linspace", "logspace"])
 @pytest.mark.parametrize(
     "start_dtype", [numpy.float64, numpy.float32, numpy.int64, numpy.int32]
@@ -778,57 +874,6 @@ def test_space_numpy_dtype(func, start_dtype, stop_dtype):
     getattr(dpnp, func)(start, stop, 10)
 
 
-@pytest.mark.parametrize(
-    "start",
-    [
-        dpnp.array(1),
-        dpnp.array([2.6]),
-        numpy.array([[-6.7, 3]]),
-        [1, -4],
-        (3, 5),
-    ],
-)
-@pytest.mark.parametrize(
-    "stop",
-    [
-        dpnp.array([-4]),
-        dpnp.array([[2.6], [-4]]),
-        numpy.array(2),
-        [[-4.6]],
-        (3,),
-    ],
-)
-def test_linspace_arrays(start, stop):
-    func = lambda xp: xp.linspace(get_array(xp, start), get_array(xp, stop), 10)
-    assert func(numpy).shape == func(dpnp).shape
-
-
-def test_linspace_complex():
-    func = lambda xp: xp.linspace(0, 3 + 2j, num=1000)
-    assert_allclose(func(dpnp), func(numpy))
-
-
-@pytest.mark.parametrize("axis", [0, 1])
-def test_linspace_axis(axis):
-    func = lambda xp: xp.linspace([2, 3], [20, 15], num=10, axis=axis)
-    assert_allclose(func(dpnp), func(numpy))
-
-
-def test_linspace_step_nan():
-    func = lambda xp: xp.linspace(1, 2, num=0, endpoint=False)
-    assert_allclose(func(dpnp), func(numpy))
-
-
-@pytest.mark.parametrize("start", [1, [1, 1]])
-@pytest.mark.parametrize("stop", [10, [10 + 10]])
-def test_linspace_retstep(start, stop):
-    func = lambda xp: xp.linspace(start, stop, num=10, retstep=True)
-    np_res = func(numpy)
-    dpnp_res = func(dpnp)
-    assert_allclose(dpnp_res[0], np_res[0])
-    assert_allclose(dpnp_res[1], np_res[1])
-
-
 @pytest.mark.parametrize(
     "arrays",
     [[], [[1]], [[1, 2, 3], [4, 5, 6]], [[1, 2], [3, 4], [5, 6]]],
@@ -862,10 +907,8 @@ def test_geomspace_zero_error():
 
 def test_space_num_error():
     with pytest.raises(ValueError):
-        dpnp.linspace(2, 5, -3)
         dpnp.geomspace(2, 5, -3)
         dpnp.logspace(2, 5, -3)
-        dpnp.linspace([2, 3], 5, -3)
         dpnp.geomspace([2, 3], 5, -3)
         dpnp.logspace([2, 3], 5, -3)
 

dpnp/tests/test_arraypad.py

@@ -73,11 +73,8 @@ def test_non_contiguous_array(self, mode):
         else:
             assert_array_equal(result, expected)
 
-    # TODO: include "linear_ramp" when dpnp issue gh-2084 is resolved
     @pytest.mark.parametrize("pad_width", [0, (0, 0), ((0, 0), (0, 0))])
-    @pytest.mark.parametrize(
-        "mode", [m for m in _modes if m not in {"linear_ramp"}]
-    )
+    @pytest.mark.parametrize("mode", _modes)
     def test_zero_pad_width(self, pad_width, mode):
         arr = dpnp.arange(30).reshape(6, 5)
         assert_array_equal(arr, dpnp.pad(arr, pad_width, mode=mode))

dpnp/tests/third_party/cupy/creation_tests/test_ranges.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import functools
 import math
 import unittest
@@ -226,8 +228,8 @@ def test_linspace_mixed_start_stop2(self, xp, dtype_range, dtype_out):
         # TODO (ev-br): np 2.0: had to bump the default rtol on Windows
         #               and numpy 1.26+weak promotion from 0 to 5e-6
         if xp.dtype(dtype_range).kind == "u":
-            # to avoid overflow, limit `val` to be smaller
-            # than xp.iinfo(dtype).max
+            # to avoid overflow, limit `val` to be smaller than
+            # xp.iinfo(dtype).max (TODO: check if dpctl-2230 resolves that)
             if dtype_range in [xp.uint8, xp.uint16] or dtype_out in [
                 xp.int8,
                 xp.uint8,

dpnp/tests/third_party/cupy/functional_tests/test_piecewise.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import unittest
 
 import numpy