Repository: cupy/cupy
Pull request: support array-like start/stop and add axis argument to linspace
This supports array-like inputs and an `axis` argument in linspace as implemented in NumPy 1.16+.

closes #2446 

~I have marked this as WIP due to what is currently relatively slow performance for the array-like code path.~ (**edit:** slow for small `num`. For large enough `num`, cupy should always be faster)

A basic translation of the code from numpy 1.17 ends up being pretty slow for small arrays. It can be faster to just run `numpy.linspace` and transfer the result to the GPU in the case of array-like `start` and/or `stop` (as currently done in #2399).

For now, I check if both start & stop are scalars and run the current implementation as a fast path. This preserves performance for this common case, but results in some code duplication.


Problem statement:
Add support for array-like inputs for linspace
Currently, numpy supports array-like inputs while cupy only considers scalars for start & stop values of linspace function. It would be great to match the functionality. 

Sample code to reproduce the feature:

```
import numpy as np
a = np.linspace(0, 9, 10)
b = np.linspace(10, 19, 10)
result = np.linspace(a, b, 10)

#####################

import cupy as cp
a = cp.linspace(0, 9, 10)
b = cp.linspace(10, 19, 10)
result = cp.linspace(a, b, 10)

```

Hints from issue discussion:

Commit message: add linspace tests with array start/stop and axis

Patch to review:
diff --git a/cupy/creation/ranges.py b/cupy/creation/ranges.py
index ca30cda62d0..2009bde7917 100644
--- a/cupy/creation/ranges.py
+++ b/cupy/creation/ranges.py
@@ -58,7 +58,8 @@ def arange(start, stop=None, step=1, dtype=None):
     return ret
 
 
-def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None):
+def _linspace_scalar(start, stop, num=50, endpoint=True, retstep=False,
+                     dtype=None):
     """Returns an array with evenly-spaced values within a given interval.
 
     Instead of specifying the step width like :func:`cupy.arange`, this
@@ -79,9 +80,6 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None):
         cupy.ndarray: The 1-D array of ranged values.
 
     """
-    if num < 0:
-        raise ValueError('linspace with num<0 is not supported')
-
     if dtype is None:
         # In actual implementation, only float is used
         dtype = float
@@ -113,6 +111,89 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None):
         return ret
 
 
+def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None,
+             axis=0):
+    """Returns an array with evenly-spaced values within a given interval.
+
+    Instead of specifying the step width like :func:`cupy.arange`, this
+    function requires the total number of elements specified.
+
+    Args:
+        start (scalar or array_like): Starting value(s) of the sequence.
+        stop (scalar or array_like): Ending value(s) of the sequence, unless
+            `endpoint` is set to False. In that case, the sequence consists of
+            all but the last of ``num + 1`` evenly spaced samples, so that
+            `stop` is excluded.  Note that the step size changes when
+            `endpoint` is False.
+        num: Number of elements.
+        endpoint (bool): If ``True``, the stop value is included as the last
+            element. Otherwise, the stop value is omitted.
+        retstep (bool): If ``True``, this function returns (array, step).
+            Otherwise, it returns only the array.
+        dtype: Data type specifier. It is inferred from the start and stop
+            arguments by default.
+        axis (int):  The axis in the result to store the samples.  Relevant
+            only if start or stop are array-like.  By default (0), the samples
+            will be along a new axis inserted at the beginning. Use -1 to get
+            an axis at the end.
+
+    Returns:
+        cupy.ndarray: The 1-D array of ranged values.
+
+    """
+    if num < 0:
+        raise ValueError('linspace with num<0 is not supported')
+    div = (num - 1) if endpoint else num
+
+    if cupy.isscalar(start) and cupy.isscalar(stop):
+        return _linspace_scalar(start, stop, num, endpoint, retstep, dtype)
+
+    start = cupy.asarray(start) * 1.0
+    stop = cupy.asarray(stop) * 1.0
+
+    dt = cupy.result_type(start, stop, float(num))
+    if dtype is None:
+        # In actual implementation, only float is used
+        dtype = dt
+
+    delta = stop - start
+
+    # ret = cupy.arange(0, num, dtype=dt).reshape((-1,) + (1,) * delta.ndim)
+    ret = cupy.empty((num,), dtype=dt)
+    _arange_ufunc(0.0, 1.0, ret, dtype=dt)
+    ret = ret.reshape((-1,) + (1,) * delta.ndim)
+
+    # In-place multiplication y *= delta/div is faster, but prevents the
+    # multiplicant from overriding what class is produced, and thus prevents,
+    # e.g. use of Quantities, see numpy#7142. Hence, we multiply in place only
+    # for standard scalar types.
+    if num > 1:
+        step = delta / div
+        if cupy.any(step == 0):
+            # Special handling for denormal numbers, numpy#5437
+            ret /= div
+            ret = ret * delta
+        else:
+            ret = ret * step
+    else:
+        # 0 and 1 item long sequences have an undefined step
+        step = float('nan')
+        # Multiply with delta to allow possible override of output class.
+        ret = ret * delta
+
+    ret += start
+    if endpoint and num > 1:
+        ret[-1] = stop
+
+    if axis != 0:
+        ret = cupy.moveaxis(ret, 0, axis)
+
+    if retstep:
+        return ret.astype(dtype, copy=False), step
+    else:
+        return ret.astype(dtype, copy=False)
+
+
 def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None):
     """Returns an array with evenly-spaced values on a log-scale.
 
@@ -318,7 +399,6 @@ def __len__(self):
      ('DD->D', 'out0 = in0 + double(i) * in1')),
     'out0 = in0 + i * in1')
 
-
 _linspace_ufunc = core.create_ufunc(
     'cupy_linspace',
     ('dd->d',),
diff --git a/tests/cupy_tests/creation_tests/test_ranges.py b/tests/cupy_tests/creation_tests/test_ranges.py
index 527c3671d9f..e97ef58fc89 100644
--- a/tests/cupy_tests/creation_tests/test_ranges.py
+++ b/tests/cupy_tests/creation_tests/test_ranges.py
@@ -145,6 +145,34 @@ def test_linspace_float_underflow(self, xp):
             x /= 2
         return xp.linspace(0., x, 10, dtype=float)
 
+    @testing.with_requires('numpy>=1.16')
+    @testing.numpy_cupy_array_equal()
+    def test_linspace_array_start_stop(self, xp):
+        start = xp.array([-120, 120], dtype="int8")
+        stop = xp.array([100, -100], dtype="int8")
+        return xp.linspace(start, stop, num=50, dtype=float)
+
+    @testing.with_requires('numpy>=1.16')
+    @testing.numpy_cupy_array_equal()
+    def test_linspace_mixed_start_stop(self, xp):
+        start = 0.0
+        stop = xp.array([100, -100], dtype="int8")
+        return xp.linspace(start, stop, num=50)
+
+    @testing.with_requires('numpy>=1.16')
+    @testing.numpy_cupy_array_equal()
+    def test_linspace_mixed_start_stop2(self, xp):
+        start = xp.array([-120, 120], dtype="int8")
+        stop = 0
+        return xp.linspace(start, stop, num=50)
+
+    @testing.with_requires('numpy>=1.16')
+    @testing.numpy_cupy_array_equal()
+    def test_linspace_array_start_stop_axis1(self, xp):
+        start = xp.array([-120, 120], dtype="int8")
+        stop = xp.array([100, -100], dtype="int8")
+        return xp.linspace(start, stop, num=50, dtype=float, axis=1)
+
     @testing.for_all_dtypes(no_bool=True)
     @testing.numpy_cupy_allclose()
     def test_logspace(self, xp, dtype):


Code review comments:
diff_hunk: @@ -145,6 +145,34 @@ def test_linspace_float_underflow(self, xp):
             x /= 2
         return xp.linspace(0., x, 10, dtype=float)
 
+    @testing.with_requires('numpy>=1.16')
+    @testing.numpy_cupy_array_equal()
+    def test_linspace_array_start_stop(self, xp):
+        start = xp.array([-120, 120], dtype="int8")
+        stop = xp.array([100, -100], dtype="int8")
+        return xp.linspace(start, stop, num=50, dtype=float)
+
+    @testing.with_requires('numpy>=1.16')
+    @testing.numpy_cupy_array_equal()
+    def test_linspace_mixed_start_stop(self, xp):
+        start = 0.0
+        stop = xp.array([100, -100], dtype="int8")
path: tests/cupy_tests/creation_tests/test_ranges.py
text: @user1:
Could you parameterize dtype with `for_all_dtypes` or `for_all_dtypes_combination`? This test seems to fail if float32 array is given as the input.

@author:

Okay, I have updated with more comprehensive testing.

There was indeed a subtle internal difference in the in dtypes being computed internally due to  #3035

If #3035 is addressed, then the refactoring done in 5dc9b9d should not be necessary.