Add support for repeat (#278)

* add repeat function

* fix styling

* optimizing repeat

* fixed minor issues

* not sure why that folder is there xD

* fixed now for sure

* test repeat not repeat test

* Fixed

---------

Co-authored-by: Bahaa Eddin tabbakha <bahaa@Bahaas-MacBook-Pro.local>

docs/src/python/ops.rst

@@ -77,6 +77,7 @@ Operations
    quantize
    quantized_matmul
    reciprocal
+   repeat
    reshape
    round
    rsqrt

mlx/ops.cpp

@@ -718,6 +718,40 @@ array stack(const std::vector<array>& arrays, StreamOrDevice s /* = {} */) {
   return stack(arrays, 0, s);
 }
 
+/** array repeat with axis */
+array repeat(const array& arr, int repeats, int axis, StreamOrDevice s) {
+  axis = normalize_axis(axis, arr.ndim());
+
+  if (repeats < 0) {
+    throw std::invalid_argument(
+        "[repeat] Number of repeats cannot be negative");
+  }
+
+  if (repeats == 0) {
+    return array({}, arr.dtype());
+  }
+
+  if (repeats == 1) {
+    return arr;
+  }
+
+  std::vector<int> new_shape(arr.shape());
+  new_shape[axis] *= repeats;
+
+  std::vector<array> repeated_arrays;
+  repeated_arrays.reserve(repeats);
+
+  for (int i = 0; i < repeats; ++i) {
+    repeated_arrays.push_back(expand_dims(arr, -1, s));
+  }
+  array repeated = concatenate(repeated_arrays, axis + 1, s);
+  return reshape(repeated, new_shape, s);
+}
+
+array repeat(const array& arr, int repeats, StreamOrDevice s) {
+  return repeat(flatten(arr, s), repeats, 0, s);
+}
+
 /** Pad an array with a constant value */
 array pad(
     const array& a,

mlx/ops.h

@@ -214,6 +214,10 @@ array concatenate(const std::vector<array>& arrays, StreamOrDevice s = {});
 array stack(const std::vector<array>& arrays, int axis, StreamOrDevice s = {});
 array stack(const std::vector<array>& arrays, StreamOrDevice s = {});
 
+/** Repeate an array along an axis. */
+array repeat(const array& arr, int repeats, int axis, StreamOrDevice s = {});
+array repeat(const array& arr, int repeats, StreamOrDevice s = {});
+
 /** Permutes the dimensions according to the given axes. */
 array transpose(const array& a, std::vector<int> axes, StreamOrDevice s = {});
 inline array transpose(

python/src/ops.cpp

@@ -2405,6 +2405,40 @@ void init_ops(py::module_& m) {
       Returns:
           array: The resulting stacked array.
     )pbdoc");
+  m.def(
+      "repeat",
+      [](const array& array,
+         int repeats,
+         std::optional<int> axis,
+         StreamOrDevice s) {
+        if (axis.has_value()) {
+          return repeat(array, repeats, axis.value(), s);
+        } else {
+          return repeat(array, repeats, s);
+        }
+      },
+      "array"_a,
+      py::pos_only(),
+      "repeats"_a,
+      "axis"_a = none,
+      py::kw_only(),
+      "stream"_a = none,
+      R"pbdoc(
+      repeat(array: array, repeats: int, axis: Optional[int] = None, *, stream: Union[None, Stream, Device] = None) -> array
+
+      Repeate an array along a specified axis.
+
+      Args:
+          array (array): Input array.
+          repeats (int): The number of repetitions for each element.
+          axis (int, optional): The axis in which to repeat the array along. If
+            unspecified it uses the flattened array of the input and repeates 
+            along axis 0.
+          stream (Stream, optional): Stream or device. Defaults to ``None``.
+
+      Returns:
+          array: The resulting repeated array.
+    )pbdoc");
   m.def(
       "clip",
       [](const array& a,

python/tests/test_ops.py

@@ -1516,6 +1516,37 @@ class TestOps(mlx_tests.MLXTestCase):
         expected = mx.array(np.linspace(0, 1, 10))
         self.assertEqualArray(d, expected)
 
+    def test_repeat(self):
+        # Setup data for the tests
+        data = np.array([[[13, 3], [16, 6]], [[14, 4], [15, 5]], [[11, 1], [12, 2]]])
+        # Test repeat along axis 0
+        repeat_axis_0 = mx.repeat(mx.array(data), 2, axis=0)
+        expected_axis_0 = np.repeat(data, 2, axis=0)
+
+        self.assertEqualArray(repeat_axis_0, mx.array(expected_axis_0))
+
+        # Test repeat along axis 1
+        repeat_axis_1 = mx.repeat(mx.array(data), 2, axis=1)
+        expected_axis_1 = np.repeat(data, 2, axis=1)
+        self.assertEqualArray(repeat_axis_1, mx.array(expected_axis_1))
+
+        # Test repeat along the last axis (default)
+        repeat_axis_2 = mx.repeat(mx.array(data), 2)
+        expected_axis_2 = np.repeat(data, 2)
+        self.assertEqualArray(repeat_axis_2, mx.array(expected_axis_2))
+
+        # Test repeat with a 1D array along axis 0
+        data_2 = mx.array([1, 3, 2])
+        repeat_2 = mx.repeat(mx.array(data_2), 3, axis=0)
+        expected_2 = np.repeat(data_2, 3)
+        self.assertEqualArray(repeat_2, mx.array(expected_2))
+
+        # Test repeat with a 2D array along axis 0
+        data_3 = mx.array([[1, 2, 3], [4, 5, 4], [0, 1, 2]])
+        repeat_3 = mx.repeat(mx.array(data_3), 2, axis=0)
+        expected_3 = np.repeat(data_3, 2, axis=0)
+        self.assertEqualArray(repeat_3, mx.array(expected_3))
+
 
 if __name__ == "__main__":
     unittest.main()

tests/ops_tests.cpp

@@ -2233,3 +2233,48 @@ TEST_CASE("test quantize dequantize") {
     CHECK(max_diff <= 127.0 / (1 << i));
   }
 }
+
+TEST_CASE("test repeat") {
+  auto data = array({13, 3, 16, 6, 14, 4, 15, 5, 11, 1, 12, 2}, {3, 2, 2});
+  auto repeat_axis_0 = repeat(data, 2, 0);
+  auto expected_axis_0 = array(
+      {13, 3, 16, 6, 13, 3, 16, 6, 14, 4, 15, 5,
+       14, 4, 15, 5, 11, 1, 12, 2, 11, 1, 12, 2},
+      {6, 2, 2});
+
+  auto repeat_axis_1 = repeat(data, 2, 1);
+  auto expected_axis_1 = array(
+      {13, 3, 13, 3, 16, 6, 16, 6, 14, 4, 14, 4,
+       15, 5, 15, 5, 11, 1, 11, 1, 12, 2, 12, 2},
+      {3, 4, 2});
+
+  auto repeat_axis_2 = repeat(data, 2); // default axis == ndim - 1 == 2
+  auto expected_axis_2 = array(
+      {13, 13, 3, 3, 16, 16, 6, 6, 14, 14, 4, 4,
+       15, 15, 5, 5, 11, 11, 1, 1, 12, 12, 2, 2},
+      {24});
+
+  // check output
+  CHECK(array_equal(repeat_axis_0, expected_axis_0).item<bool>());
+  CHECK(array_equal(repeat_axis_1, expected_axis_1).item<bool>());
+  CHECK(array_equal(repeat_axis_2, expected_axis_2).item<bool>());
+
+  auto data_2 = array({1, 3, 2}, {3});
+  auto repeat_2 = repeat(data_2, 2, 0);
+  auto expected_2 = array({1, 1, 3, 3, 2, 2}, {6});
+  CHECK(array_equal(repeat_2, expected_2).item<bool>());
+
+  auto data_3 = array({1, 2, 3, 4, 5, 4, 0, 1, 2}, {3, 3});
+  auto repeat_3 = repeat(data_3, 2, 0);
+  auto expected_3 =
+      array({1, 2, 3, 1, 2, 3, 4, 5, 4, 4, 5, 4, 0, 1, 2, 0, 1, 2}, {6, 3});
+  CHECK(array_equal(repeat_3, expected_3).item<bool>());
+
+  // 0 repeats
+  auto repeat_4 = repeat(data_3, 0, 0);
+  auto expected_4 = array({});
+  CHECK(array_equal(repeat_2, expected_2).item<bool>());
+
+  // negative repeats
+  CHECK_THROWS_AS(repeat(data_3, -3, 0), std::invalid_argument);
+}