(Re)implement ndonnx.repeat for all built-in data types (#161)

Author: cbourjau

CHANGELOG.rst

@@ -7,9 +7,13 @@
 Changelog
 =========
 
-0.14.1 (unreleased)
+0.15.0 (2025-08-13)
 -------------------
 
+**New feature**
+
+- :func:`ndonnx.repeat` is now implemented for all built-in data types.
+
 **Other changes**
 
 - Improve stacklevel presented of various deprecation warnings.

array-api-tests

@@ -840,11 +840,6 @@ def reshape(self, shape: tuple[int, ...] | TyArrayInt64) -> Self:
             var = op.reshape(self._var, shape._var, allowzero=True)
         return type(self)(var)
 
-    def repeat(
-        self, repeats: int | TyArrayInt64, /, *, axis: int | None = None
-    ) -> Self:
-        raise NotImplementedError
-
     def searchsorted(
         self,
         x2: Self,
@@ -2820,3 +2815,18 @@ def dummy_n_m(a: TY_ARRAY_NUMBER, b: TY_ARRAY_NUMBER) -> TyArrayInt64:
 
     arr = type(a)(var_with_propagation)
     return arr
+
+
+def arange(
+    start: TyArrayInt64,
+    /,
+    stop: TyArrayInt64 | None = None,
+    step: int | TyArrayInt64 = 1,
+) -> TyArrayInt64:
+    if stop is None:
+        stop = start
+        start = const(0)
+    if isinstance(step, int):
+        step = const(step)
+    var = op.range(start=start._var, limit=stop._var, delta=step._var)
+    return TyArrayInt64(var)

ndonnx/_typed_array/onnx.py

@@ -457,7 +457,58 @@ def variance(
     def repeat(
         self, repeats: int | TyArrayInt64, /, *, axis: int | None = None
     ) -> Self:
-        raise _make_type_error("repeat", self.dtype)
+        # Default implementation of `repeat` that does not depend on the underlying layout
+        from . import onnx
+
+        x = self
+        if axis is None:
+            x = x.reshape((-1,))
+            axis = 0
+
+        if axis < 0:
+            axis += x.ndim
+
+        x_shape = x.dynamic_shape
+
+        if isinstance(repeats, onnx.TyArrayInt64) and repeats.ndim > 1:
+            raise ValueError(
+                f"'repeats' must be 0 or 1 dimensional, but has `{repeats.ndim}` dimensions"
+            )
+
+        if isinstance(repeats, int):
+            repeats = onnx.const(np.asarray(repeats), onnx.int64)
+
+        if repeats.ndim == 0:
+            repeats = repeats[None]
+            # Expand by one dimension and use regular broadcasting to repeat
+            key: list[slice | None] = [slice(None, None)] * (x.ndim + 1)
+            key[axis + 1] = None
+            tmp = x[tuple(key)].broadcast_to(
+                x_shape[: axis + 1].concat([repeats, x_shape[axis + 1 :]], axis=0)
+            )
+            # We now have a shape of (..., self.shape[axis], repeats, ...).
+            # All we need to do is collapse these two axes into one.
+            out_shape = x_shape.copy()
+            out_shape[axis] = safe_cast(onnx.TyArrayInt64, out_shape[axis] * repeats)
+
+            return tmp.reshape(out_shape)
+
+        # Repeats may be of shape (1,)
+        repeats = repeats.broadcast_to(x_shape[axis][None])
+
+        max_rep = repeats.max(keepdims=False)
+        repeated_with_maxrep = x.repeat(max_rep, axis=axis)
+
+        mask2D = onnx.arange(max_rep)[None, :] < repeats[:, None]
+        mask1D = mask2D.reshape((-1,))
+
+        # We can only use __getitem__ with a boolean mask on 1D arrays.
+        # Otherwise, we have to compute indices and use `take`
+        if x.ndim == 1:
+            return repeated_with_maxrep[mask1D]
+
+        indices = onnx.arange(mask1D.dynamic_size)[mask1D]
+        return repeated_with_maxrep.take(indices, axis=axis)
 
     def tile(self, repetitions: tuple[int, ...], /) -> Self:
         raise _make_type_error("tile", self.dtype)

ndonnx/_typed_array/typed_array.py

@@ -1,6 +1,5 @@
 # Missing implementations which should be perfectly doable
 array_api_tests/test_linalg.py::test_tensordot
-array_api_tests/test_manipulation_functions.py::test_repeat
 
 # Tests that fail due to lossy ORT workarounds, which are tested elsewhere
 array_api_tests/test_statistical_functions.py::test_sum

skips.txt

@@ -15,3 +15,39 @@ def do(npx):
         return npx.concat([a1, a2], axis=None)
 
     np.testing.assert_array_equal(do(np), do(ndx))
+
+
+@pytest.mark.parametrize(
+    "x, repeats, axis",
+    [
+        ([1, 2, 3], 2, None),
+        ([1, 2, 3], np.asarray(2), None),
+        ([1, 2, 3], np.asarray(2), 0),
+        ([1, 2, 3], np.asarray([2]), None),
+        ([1, 2, 3], [1, 0, 2], None),
+        ([1, 2, 3], [1, 0, 2], 0),
+        ([[1, 2, 3]], [1, 0, 2], None),
+        ([[1, 2, 3]], [1, 0, 2], 1),
+        (np.arange(27).reshape((3, 3, 3)), [1, 0, 2], 0),
+        (np.arange(27).reshape((3, 3, 3)), [1, 0, 2], 1),
+        (np.arange(27).reshape((3, 3, 3)), [1, 0, 2], 2),
+        # zero-size
+        (np.ones((3, 0, 3)), 2, None),
+        (np.ones((3, 0, 3)), 2, 0),
+        (np.ones((3, 0, 3)), 2, 1),
+        (np.ones((3, 0, 3)), 2, 2),
+        (np.ones((3, 0, 3)), [2], 2),
+        (np.ones((3, 0, 3)), [1, 0, 2], 2),
+        # other data types
+        (["a", "b", "c"], np.asarray([2]), None),
+        (np.asarray([1, 2, 3], dtype="datetime64[s]"), np.asarray([2]), None),
+    ],
+)
+def test_repeat(x, repeats, axis):
+    def do(npx):
+        repeats_ = repeats
+        if isinstance(repeats, list | np.ndarray):
+            repeats_ = npx.asarray(repeats_)
+        return npx.repeat(npx.asarray(x), repeats_, axis=axis)
+
+    np.testing.assert_array_equal(do(ndx).unwrap_numpy(), do(np))