Implements dpctl.tensor.repeat

Author: ndgrigorian

dpctl/tensor/CMakeLists.txt

@@ -49,6 +49,7 @@ pybind11_add_module(${python_module_name} MODULE
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/device_support_queries.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/elementwise_functions.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/sum_reductions.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/repeat.cpp
 )
 set(_clang_prefix "")
 if (WIN32)

dpctl/tensor/__init__.py

@@ -72,6 +72,7 @@
     iinfo,
     moveaxis,
     permute_dims,
+    repeat,
     result_type,
     roll,
     squeeze,
@@ -305,4 +306,5 @@
     "tanh",
     "trunc",
     "allclose",
+    "repeat",
 ]

dpctl/tensor/_manipulation_functions.py

@@ -15,9 +15,12 @@
 #  limitations under the License.
 
 
-from itertools import chain, product, repeat
+import itertools
+import operator
+from itertools import chain, product
 
 import numpy as np
+from numpy import AxisError
 from numpy.core.numeric import normalize_axis_index, normalize_axis_tuple
 
 import dpctl
@@ -132,7 +135,7 @@ def _broadcast_shape_impl(shapes):
         diff = biggest - nds[i]
         if diff > 0:
             ty = type(shapes[i])
-            shapes[i] = ty(chain(repeat(1, diff), shapes[i]))
+            shapes[i] = ty(chain(itertools.repeat(1, diff), shapes[i]))
     common_shape = []
     for axis in range(biggest):
         lengths = [s[axis] for s in shapes]
@@ -928,6 +931,94 @@ def swapaxes(X, axis1, axis2):
     return dpt.permute_dims(X, tuple(ind))
 
 
+def repeat(x, repeats, axis=None):
+    if not isinstance(x, dpt.usm_ndarray):
+        raise TypeError(f"Expected usm_ndarray type, got {type(x)}.")
+
+    x_ndim = x.ndim
+    if axis is None:
+        if x_ndim != 1:
+            raise ValueError(
+                f"`axis` cannot be `None` for array of dimension {x_ndim}"
+            )
+        axis = 0
+
+    x_shape = x.shape
+    if x_ndim > 0:
+        axis = normalize_axis_index(operator.index(axis), x_ndim)
+        axis_size = x_shape[axis]
+    else:
+        if axis != 0:
+            AxisError("`axis` must be `0` for input of dimension `0`")
+        axis_size = x.size
+
+    scalar = False
+    if isinstance(repeats, int):
+        scalar = True
+        if repeats < 0:
+            raise ValueError("`repeats` must be a positive integer")
+    elif isinstance(repeats, tuple):
+        len_reps = len(repeats)
+        if len_reps != axis_size:
+            raise ValueError(
+                "`repeats` tuple must have the same length as the repeated axis"
+            )
+        elif len_reps == 1:
+            repeats = repeats[0]
+            if repeats < 0:
+                raise ValueError("`repeats` elements must be positive")
+            scalar = True
+    else:
+        raise TypeError(
+            f"Expected int or tuple for second argument, got {type(repeats)}"
+        )
+
+    usm_type = x.usm_type
+    exec_q = x.sycl_queue
+
+    if axis_size == 0:
+        return dpt.empty(x_shape, dtype=x.dtype, sycl_queue=exec_q)
+
+    if scalar:
+        res_axis_size = repeats * axis_size
+        res_shape = x_shape[:axis] + (res_axis_size,) + x_shape[axis + 1 :]
+        res = dpt.empty(
+            res_shape, dtype=x.dtype, usm_type=usm_type, sycl_queue=exec_q
+        )
+        if res_axis_size > 0:
+            hev, _ = ti._repeat_by_scalar(
+                src=x,
+                dst=res,
+                reps=repeats,
+                axis=axis,
+                sycl_queue=exec_q,
+            )
+            hev.wait()
+    else:
+        repeats = dpt.asarray(repeats, dtype="i8", sycl_queue=exec_q)
+        if not dpt.all(repeats >= 0):
+            raise ValueError("`repeats` elements must be positive")
+        cumsum = dpt.empty(
+            (axis_size,), dtype=dpt.int64, usm_type=usm_type, sycl_queue=exec_q
+        )
+        res_axis_size = ti._cumsum_1d(repeats, cumsum, sycl_queue=exec_q)
+        res_shape = x_shape[:axis] + (res_axis_size,) + x_shape[axis + 1 :]
+        res = dpt.empty(
+            res_shape, dtype=x.dtype, usm_type=usm_type, sycl_queue=exec_q
+        )
+        if res_axis_size > 0:
+            hev, _ = ti._repeat_by_sequence(
+                src=x,
+                dst=res,
+                reps=repeats,
+                cumsum=cumsum,
+                axis=axis,
+                sycl_queue=exec_q,
+            )
+            hev.wait()
+    return res
+
+
 def _supported_dtype(dtypes):
     for dtype in dtypes:
         if dtype.char not in "?bBhHiIlLqQefdFD":