Merge pull request #855 from IntelPython/bugfix-arange

Fixed bug causing dpt.arange(0,1,1e-3,dtype="f4") be length 1

Author: oleksandr-pavlyk

dpctl/tensor/_ctors.py

@@ -471,22 +471,24 @@ def _coerce_and_infer_dt(*args, dt):
         raise ValueError(f"Data type {dt} is not supported")
 
 
+def _round_for_arange(tmp):
+    k = int(tmp)
+    if k > 0 and float(k) < tmp:
+        tmp = tmp + 1
+    return tmp
+
+
 def _get_arange_length(start, stop, step):
     "Compute length of arange sequence"
     span = stop - start
     if type(step) in [int, float] and type(span) in [int, float]:
-        offset = -1 if step > 0 else 1
-        tmp = 1 + (span + offset) / step
-        return tmp
+        return _round_for_arange(span / step)
     tmp = span / step
     if type(tmp) is complex and tmp.imag == 0:
         tmp = tmp.real
     else:
         return tmp
-    k = int(tmp)
-    if k > 0 and float(k) < tmp:
-        tmp = tmp + 1
-    return tmp
+    return _round_for_arange(tmp)
 
 
 def arange(

dpctl/tests/test_usm_ndarray_ctor.py

@@ -994,6 +994,20 @@ def test_arange(dt):
     assert X2.shape == (sz,)
 
 
+def test_arange_fp():
+    try:
+        q = dpctl.SyclQueue()
+    except dpctl.SyclQueueCreationError:
+        pytest.skip("Queue could not be created")
+
+    assert dpt.arange(7, 0, -2, dtype="f4", device=q).shape == (4,)
+    assert dpt.arange(0, 1, 0.25, dtype="f4", device=q).shape == (4,)
+
+    if q.sycl_device.has_aspect_fp64:
+        assert dpt.arange(7, 0, -2, dtype="f8", device=q).shape == (4,)
+    assert dpt.arange(0, 1, 0.25, dtype="f4", device=q).shape == (4,)
+
+
 @pytest.mark.parametrize(
     "dt",
     _all_dtypes,