Add repeat to the specification

data-apis · rgommers · Feb 22, 2024 · Sep 21, 2023 · Sep 21, 2023 · Sep 21, 2023
commit 8764700533a0308a065d8f5024121ef8ef0ee3a9
diff --git a/spec/draft/API_specification/manipulation_functions.rst b/spec/draft/API_specification/manipulation_functions.rst
@@ -25,6 +25,7 @@ Objects in API
    flip
    moveaxis
    permute_dims
+   repeat
    reshape
    roll
    squeeze

diff --git a/src/array_api_stubs/_draft/manipulation_functions.py b/src/array_api_stubs/_draft/manipulation_functions.py
@@ -6,6 +6,7 @@
     "flip",
     "moveaxis",
     "permute_dims",
+    "repeat",
     "reshape",
     "roll",
     "squeeze",
@@ -14,7 +15,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]:
@@ -158,6 +159,46 @@ 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.
+    axis: Optional[int]
+        the axis (dimension) along which to repeat elements. If ``axis`` is `None`, the function must repeat elements of a flattened input array ``x`` and return the result as a one-dimensional output array. 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:
-Original file line number
+Diff line change
@@ Expand Up / @@ -25,6 +25,7 @@ Objects in API @@
        flip
        moveaxis
        permute_dims
+       repeat
        reshape
        roll
        squeeze
@@ Expand Down @@