From ff2b58e299df423f0dd6520deb82ee7e04b77a53 Mon Sep 17 00:00:00 2001 From: Bahaa <60551967+Bahaatbb@users.noreply.github.com> Date: Thu, 28 Dec 2023 00:11:38 +0300 Subject: [PATCH] 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 --- docs/src/python/ops.rst | 1 + mlx/ops.cpp | 34 ++++++++++++++++++++++++++++++ mlx/ops.h | 4 ++++ python/src/ops.cpp | 34 ++++++++++++++++++++++++++++++ python/tests/test_ops.py | 31 +++++++++++++++++++++++++++ tests/ops_tests.cpp | 45 ++++++++++++++++++++++++++++++++++++++++ 6 files changed, 149 insertions(+) diff --git a/docs/src/python/ops.rst b/docs/src/python/ops.rst index 0c5763290..6c31b54ec 100644 --- a/docs/src/python/ops.rst +++ b/docs/src/python/ops.rst @@ -77,6 +77,7 @@ Operations quantize quantized_matmul reciprocal + repeat reshape round rsqrt diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 4a7f39321..eac77735b 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -718,6 +718,40 @@ array stack(const std::vector& 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 new_shape(arr.shape()); + new_shape[axis] *= repeats; + + std::vector 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, diff --git a/mlx/ops.h b/mlx/ops.h index e1abac6fb..a99465f3e 100644 --- a/mlx/ops.h +++ b/mlx/ops.h @@ -214,6 +214,10 @@ array concatenate(const std::vector& arrays, StreamOrDevice s = {}); array stack(const std::vector& arrays, int axis, StreamOrDevice s = {}); array stack(const std::vector& 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 axes, StreamOrDevice s = {}); inline array transpose( diff --git a/python/src/ops.cpp b/python/src/ops.cpp index f97a55bce..8da45e0eb 100644 --- a/python/src/ops.cpp +++ b/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 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, diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index ecfc6315a..89edfdefa 100644 --- a/python/tests/test_ops.py +++ b/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() diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index 5bee4b5c8..2c6c8e8ef 100644 --- a/tests/ops_tests.cpp +++ b/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()); + CHECK(array_equal(repeat_axis_1, expected_axis_1).item()); + CHECK(array_equal(repeat_axis_2, expected_axis_2).item()); + + 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()); + + 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()); + + // 0 repeats + auto repeat_4 = repeat(data_3, 0, 0); + auto expected_4 = array({}); + CHECK(array_equal(repeat_2, expected_2).item()); + + // negative repeats + CHECK_THROWS_AS(repeat(data_3, -3, 0), std::invalid_argument); +} \ No newline at end of file