Add repeat()

Author: asmeurer

array_api_strict/__init__.py

@@ -277,13 +277,14 @@
     flip,
     moveaxis,
     permute_dims,
+    repeat,
     reshape,
     roll,
     squeeze,
     stack,
 )
 
-__all__ += ["concat", "expand_dims", "flip", "moveaxis", "permute_dims", "reshape", "roll", "squeeze", "stack"]
+__all__ += ["concat", "expand_dims", "flip", "moveaxis", "permute_dims", "repeat", "reshape", "roll", "squeeze", "stack"]
 
 from ._searching_functions import argmax, argmin, nonzero, where
 

array_api_strict/_flags.py

@@ -81,10 +81,10 @@ def set_array_api_strict_flags(
       The functions that make use of data-dependent shapes, and are therefore
       disabled by setting this flag to False are
 
-      - `unique_all`, `unique_counts`, `unique_inverse`, and `unique_values`.
-      - `nonzero`
+      - `unique_all()`, `unique_counts()`, `unique_inverse()`, and `unique_values()`.
+      - `nonzero()`
       - Boolean array indexing
-      - `repeat` when the `repeats` argument is an array (requires 2023.12
+      - `repeat()` when the `repeats` argument is an array (requires 2023.12
         version of the standard)
 
       See

array_api_strict/_manipulation_functions.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
 from ._array_object import Array
+from ._creation_functions import asarray
 from ._data_type_functions import result_type
-from ._flags import requires_api_version
+from ._flags import requires_api_version, get_array_api_strict_flags
 
 from typing import TYPE_CHECKING
 
@@ -58,7 +59,6 @@ def moveaxis(
     """
     return Array._new(np.moveaxis(x._array, source, destination))
 
-
 # Note: The function name is different here (see also matrix_transpose).
 # Unlike transpose(), the axes argument is required.
 def permute_dims(x: Array, /, axes: Tuple[int, ...]) -> Array:
@@ -69,6 +69,29 @@ def permute_dims(x: Array, /, axes: Tuple[int, ...]) -> Array:
     """
     return Array._new(np.transpose(x._array, axes))
 
+@requires_api_version('2023.12')
+def repeat(
+    x: Array,
+    repeats: Union[int, Array],
+    /,
+    *,
+    axis: Optional[int] = None,
+) -> Array:
+    """
+    Array API compatible wrapper for :py:func:`np.repeat <numpy.repeat>`.
+
+    See its docstring for more information.
+    """
+    if isinstance(repeats, Array):
+        data_dependent_shapes = get_array_api_strict_flags()['data_dependent_shapes']
+        if not data_dependent_shapes:
+            raise RuntimeError("repeat() with repeats as an array requires data-dependent shapes, but the data_dependent_shapes flag has been disabled for array-api-strict")
+    elif isinstance(repeats, int):
+        repeats = asarray(repeats)
+    else:
+        raise TypeError("repeats must be an int or array")
+
+    return Array._new(np.repeat(x._array, repeats, axis=axis))
 
 # Note: the optional argument is called 'shape', not 'newshape'
 def reshape(x: Array,

array_api_strict/tests/test_flags.py

@@ -2,7 +2,7 @@
                       reset_array_api_strict_flags)
 
 from .. import (asarray, unique_all, unique_counts, unique_inverse,
-                unique_values, nonzero)
+                unique_values, nonzero, repeat)
 
 import array_api_strict as xp
 
@@ -102,8 +102,12 @@ def test_api_version():
     assert xp.__array_api_version__ == '2021.12'
 
 def test_data_dependent_shapes():
+    with pytest.warns(UserWarning):
+        set_array_api_strict_flags(api_version='2023.12') # to enable repeat()
+
     a = asarray([0, 0, 1, 2, 2])
     mask = asarray([True, False, True, False, True])
+    repeats = asarray([1, 1, 2, 2, 2])
 
     # Should not error
     unique_all(a)
@@ -112,7 +116,8 @@ def test_data_dependent_shapes():
     unique_values(a)
     nonzero(a)
     a[mask]
-    # TODO: add repeat when it is implemented
+    repeat(a, repeats)
+    repeat(a, 2)
 
     set_array_api_strict_flags(data_dependent_shapes=False)
 
@@ -122,6 +127,8 @@ def test_data_dependent_shapes():
     pytest.raises(RuntimeError, lambda: unique_values(a))
     pytest.raises(RuntimeError, lambda: nonzero(a))
     pytest.raises(RuntimeError, lambda: a[mask])
+    pytest.raises(RuntimeError, lambda: repeat(a, repeats))
+    repeat(a, 2) # Should never error
 
 linalg_examples = {
     'cholesky': lambda: xp.linalg.cholesky(xp.eye(3)),