Add repeat to the specification

spec/draft/API_specification/manipulation_functions.rst

@@ -25,6 +25,7 @@ Objects in API
    flip
    moveaxis
    permute_dims
+   repeat
    reshape
    roll
    squeeze

src/array_api_stubs/_draft/manipulation_functions.py

@@ -6,6 +6,7 @@
     "flip",
     "moveaxis",
     "permute_dims",
+    "repeat",
     "reshape",
     "roll",
     "squeeze",
@@ -15,7 +16,7 @@
 ]
 
 
-from ._types import List, Optional, Tuple, Union, array
+from ._types import List, Optional, Tuple, Union, Sequence, array
 
 
 def broadcast_arrays(*arrays: array) -> List[array]:
@@ -159,6 +160,51 @@ def permute_dims(x: array, /, axes: Tuple[int, ...]) -> array:
     """
 
 
+def repeat(
+    x: array,
+    repeats: Union[int, Sequence[int], array],
+    /,
+    *,
+    axis: Optional[int] = None,
+) -> array:
+    """
+    Repeats elements of an array.
+
+    Parameters
+    ----------
+    x: array
+        input array containing elements to repeat.
+    repeats: Union[int, Sequence[int], array]
+        the number of repetitions for each element.
+
+        If ``axis`` is ``None``, let ``N = prod(x.shape)`` and
+
+        -   if ``repeats`` is an array, ``repeats`` must be broadcast compatible with the shape ``(N,)`` (i.e., be a one-dimensional array having shape ``(1,)`` or ``(N,)``).
+        -   if ``repeats`` is a sequence of integers, ``len(repeats)`` must be broadcast compatible with the shape ``(N,)`` (i.e., the number of sequence elements be either ``1`` or ``N``).
+        -   if ``repeats`` is an integer, ``repeats`` must be broadcasted to the shape `(N,)`.
+
+        If ``axis`` is not ``None``, let ``M = x.shape[axis]`` and
+
+        -   if ``repeats`` is an array, ``repeats`` must be broadcast compatible with the shape ``(M,)`` (i.e., be a one-dimensional array having shape ``(1,)`` or ``(M,)``).
+        -   if ``repeats`` is a sequence of integers, ``len(repeats)`` must be broadcast compatible with the shape ``(M,)`` (i.e., the number of sequence elements must be either ``1`` or ``M``).
+        -   if ``repeats`` is an integer, ``repeats`` must be broadcasted to the shape ``(M,)``.
+
+        If ``repeats`` is an array, the array must have an integer data type.
+
+
+    .. note::
+       For specification-conforming array libraries supporting hardware acceleration, providing an array for ``repeats`` may cause device synchronization due to an unknown output shape. For those array libraries where synchronization concerns are applicable, conforming array libraries are advised to include a warning in their documentation regarding potential performance degradation when ``repeats`` is an array.
+
+    axis: Optional[int]
+        the axis (dimension) along which to repeat elements. If ``axis`` is `None`, the function must flatten the input array ``x`` and then repeat elements of the flattened input array and return the result as a one-dimensional output array. A flattened input array must be flattened in row-major, C-style order. Default: ``None``.
+
+    Returns
+    -------
+    out: array
+        an output array containing repeated elements. The returned array must have the same data type as ``x``. If ``axis`` is ``None``, the returned array must be a one-dimensional array; otherwise, the returned array have the same shape as ``x``, except for the axis (dimension) along which elements were repeated.
+    """
+
+
 def reshape(
     x: array, /, shape: Tuple[int, ...], *, copy: Optional[bool] = None
 ) -> array: