Merge pull request #1504 from IntelPython/fix-copy-of-src-with-offset

Fix copy of src with offset

Author: oleksandr-pavlyk

dpctl/tensor/_copy_utils.py

@@ -300,14 +300,22 @@ def _copy_from_usm_ndarray_to_usm_ndarray(dst, src):
             src.shape, src.strides, len(common_shape)
         )
         src_same_shape = dpt.usm_ndarray(
-            common_shape, dtype=src.dtype, buffer=src, strides=new_src_strides
+            common_shape,
+            dtype=src.dtype,
+            buffer=src,
+            strides=new_src_strides,
+            offset=src._element_offset,
         )
     elif src.ndim == len(common_shape):
         new_src_strides = _broadcast_strides(
             src.shape, src.strides, len(common_shape)
         )
         src_same_shape = dpt.usm_ndarray(
-            common_shape, dtype=src.dtype, buffer=src, strides=new_src_strides
+            common_shape,
+            dtype=src.dtype,
+            buffer=src,
+            strides=new_src_strides,
+            offset=src._element_offset,
         )
     else:
         # since broadcasting succeeded, src.ndim is greater because of

dpctl/tests/test_usm_ndarray_ctor.py

@@ -1036,6 +1036,7 @@ def test_setitem_same_dtype(dtype, src_usm_type, dst_usm_type):
 
 
 def test_setitem_broadcasting():
+    "See gh-1503"
     get_queue_or_skip()
     dst = dpt.ones((2, 3, 4), dtype="u4")
     src = dpt.zeros((3, 1), dtype=dst.dtype)
@@ -1044,6 +1045,16 @@ def test_setitem_broadcasting():
     assert np.array_equal(dpt.asnumpy(dst), expected)
 
 
+def test_setitem_broadcasting_offset():
+    get_queue_or_skip()
+    dt = dpt.int32
+    x = dpt.asarray([[1, 2, 3], [6, 7, 8]], dtype=dt)
+    y = dpt.asarray([4, 5], dtype=dt)
+    x[0] = y[1]
+    expected = dpt.asarray([[5, 5, 5], [6, 7, 8]], dtype=dt)
+    assert dpt.all(x == expected)
+
+
 def test_setitem_broadcasting_empty_dst_validation():
     "Broadcasting rules apply, except exception"
     get_queue_or_skip()