Skip to content

Commit 7ed14a6

Browse files
Fixed dpt.arange(9.7, stop=200, step=100, dtype='i4')
To align with behavior of numpy, changed the logic of determining step argument for the call to `_linspace_step` routine. We now compute first and second element of the array of given type, and determine the step as a the difference of these. To avoid possible overflow message when subtracting unsigned integers, cast first and second element to int64, subtract, and cast to the target type. ```ipython In [1]: import dpctl, dpctl.tensor as dpt, numpy as np In [2]: dpt.asnumpy(dpt.arange(9.7, stop=200, step=100, dtype='i4')) Out[2]: array([ 9, 109], dtype=int32) In [3]: np.arange(9.7, stop=200, step=100, dtype='i4') Out[3]: array([ 9, 109], dtype=int32) ```
1 parent 18e8782 commit 7ed14a6

File tree

1 file changed

+12
-3
lines changed

1 file changed

+12
-3
lines changed

dpctl/tensor/_ctors.py

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -569,9 +569,18 @@ def arange(
569569
buffer_ctor_kwargs={"queue": sycl_queue},
570570
)
571571
sc_ty = dt.type
572-
_step = sc_ty(start + step) - start
573-
_step = sc_ty(_step)
574-
_start = sc_ty(start)
572+
_first = sc_ty(start)
573+
if sh > 1:
574+
_second = sc_ty(start + step)
575+
if dt in [dpt.uint8, dpt.uint16, dpt.uint32, dpt.uint64]:
576+
int64_ty = dpt.int64.type
577+
_step = int64_ty(_second) - int64_ty(_first)
578+
else:
579+
_step = _second - _first
580+
_step = sc_ty(_step)
581+
else:
582+
_step = sc_ty(1)
583+
_start = _first
575584
hev, _ = ti._linspace_step(_start, _step, res, sycl_queue)
576585
hev.wait()
577586
if is_bool:

0 commit comments

Comments
 (0)