From fc14b6a8947f197e850a7de7b66cb238b811e364 Mon Sep 17 00:00:00 2001 From: Bahaatbb Date: Sat, 30 Dec 2023 01:59:37 +0300 Subject: [PATCH 1/3] add supprot for tile --- docs/src/python/ops.rst | 1 + mlx/ops.cpp | 35 ++++++++++++++++++ mlx/ops.h | 3 ++ python/src/ops.cpp | 23 ++++++++++++ python/tests/test_ops.py | 48 ++++++++++++++++++++++++- tests/ops_tests.cpp | 77 +++++++++++++++++++++++++++++++++++++++- 6 files changed, 185 insertions(+), 2 deletions(-) diff --git a/docs/src/python/ops.rst b/docs/src/python/ops.rst index 6c31b54ec1..8c7bf011d5 100644 --- a/docs/src/python/ops.rst +++ b/docs/src/python/ops.rst @@ -104,6 +104,7 @@ Operations take_along_axis tan tanh + tile transpose tri tril diff --git a/mlx/ops.cpp b/mlx/ops.cpp index f4f6b922d5..8e153e33ac 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -753,6 +753,41 @@ array repeat(const array& arr, int repeats, StreamOrDevice s) { return repeat(flatten(arr, s), repeats, 0, s); } +/** Replicating array a specified number of times along each axis */ +array tile(const array& a, std::vector repeats, StreamOrDevice s) { + if (repeats.size() == 0) { + return a; + } + + int d = repeats.size(); + array arr = copy(a, s); + + if (d < arr.ndim()) { + repeats.insert(repeats.begin(), arr.ndim() - d, 1); + } else if (d > arr.ndim()) { + arr = expand_dims(arr, std::vector(d - arr.ndim(), 0), s); + } + + std::vector shape_out; + + for (size_t i = 0; i < arr.shape().size(); i++) { + shape_out.push_back(arr.shape()[i] * repeats[i]); + } + + int n = arr.size(); + std::vector old_shape = arr.shape(); + if (n > 0) { + for (size_t i = 0; i < old_shape.size(); i++) { + if (repeats[i] != 1) { + arr = repeat(reshape(arr, {-1, n}, s), repeats[i], 0, s); + } + n /= old_shape[i]; + } + } + + return reshape(arr, shape_out, s); +} + /** Pad an array with a constant value */ array pad( const array& a, diff --git a/mlx/ops.h b/mlx/ops.h index a99465f3e6..5b7a9a6a70 100644 --- a/mlx/ops.h +++ b/mlx/ops.h @@ -218,6 +218,9 @@ array stack(const std::vector& arrays, StreamOrDevice s = {}); array repeat(const array& arr, int repeats, int axis, StreamOrDevice s = {}); array repeat(const array& arr, int repeats, StreamOrDevice s = {}); +/** Replicating array a specified number of times along each axis */ +array tile(const array& arr, std::vector 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 8da45e0eba..a5a0726808 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -2439,6 +2439,29 @@ void init_ops(py::module_& m) { Returns: array: The resulting repeated array. )pbdoc"); + m.def( + "tile", + [](const array& array, std::vector repeats, StreamOrDevice s) { + return tile(array, repeats, s); + }, + "array"_a, + py::pos_only(), + "repeats"_a, + py::kw_only(), + "stream"_a = none, + R"pbdoc( + tile(array: array, repeats: array, *, stream: Union[None, Stream, Device] = None) -> array + + Replicating array a specified number of times along each axis. + + Args: + array (array): Input array. + repeats (array): The number of repetitions along each axis. + stream (Stream, optional): Stream or device. Defaults to ``None``. + + Returns: + array: The resulting tiled array. + )pbdoc"); m.def( "clip", [](const array& a, diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 89edfdefa1..a77359ab63 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -1518,7 +1518,10 @@ def test_linspace(self): def test_repeat(self): # Setup data for the tests - data = np.array([[[13, 3], [16, 6]], [[14, 4], [15, 5]], [[11, 1], [12, 2]]]) + data = np.array( + [[[13, 3], [16, 6]], [[14, 4], [15, 5]], [[11, 1], [12, 2]]], + dtype=np.float32, + ) # 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) @@ -1547,6 +1550,49 @@ def test_repeat(self): expected_3 = np.repeat(data_3, 2, axis=0) self.assertEqualArray(repeat_3, mx.array(expected_3)) + def test_tile(self): + data = np.array( + [ + [[6.0, 10.0], [1.0, 3.0]], + [[9.0, 5.0], [9.0, 2.0]], + [[5.0, 3.0], [3.0, 4.0]], + ], + dtype=np.float32, + ) + + expected_1 = mx.tile(mx.array(data), [3]) + tiled_1 = np.tile(data, [3]) + self.assertEqualArray(expected_1, mx.array(tiled_1)) + + data_2 = np.array( + [ + [[4.0, 4.0], [9.0, 7.0]], + [[7.0, 3.0], [6.0, 8.0]], + [[8.0, 3.0], [8.0, 7.0]], + ], + dtype=np.float32, + ) + + expected_2 = mx.tile(mx.array(data_2), [1, 3, 2]) + tiled_2 = np.tile(data_2, [1, 3, 2]) + self.assertEqualArray(expected_2, mx.array(tiled_2)) + + expected_3 = mx.tile(mx.array(data_2), [3, 2, 2]) + tiled_3 = np.tile(data_2, (3, 2, 2)) + self.assertEqualArray(expected_3, mx.array(tiled_3)) + + expected_4 = mx.tile(mx.array(data_2), [3, 2, 2, 2]) + tiled_4 = np.tile(data_2, (3, 2, 2, 2)) + self.assertEqualArray(expected_4, mx.array(tiled_4)) + + expected_5 = mx.tile(mx.array(data_2), [1, 1, 1]) + tiled_5 = np.tile(data_2, (1, 1, 1)) + self.assertEqualArray(expected_5, mx.array(tiled_5)) + + expected_6 = mx.tile(mx.array(data_2), [1, 1, 1, 1]) + tiled_6 = np.tile(data_2, (1, 1, 1, 1)) + self.assertEqualArray(expected_6, mx.array(tiled_6)) + if __name__ == "__main__": unittest.main() diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index 2c6c8e8efb..cf9a2a455f 100644 --- a/tests/ops_tests.cpp +++ b/tests/ops_tests.cpp @@ -2277,4 +2277,79 @@ TEST_CASE("test repeat") { // negative repeats CHECK_THROWS_AS(repeat(data_3, -3, 0), std::invalid_argument); -} \ No newline at end of file +} + +TEST_CASE("test tile") { + auto data = + array({6., 10., 1., 3., 9., 5., 9., 2., 5., 3., 3., 4.}, {3, 2, 2}); + auto expected_1 = array( + {6., 10., 6., 10., 6., 10., 1., 3., 1., 3., 1., 3., + 9., 5., 9., 5., 9., 5., 9., 2., 9., 2., 9., 2., + 5., 3., 5., 3., 5., 3., 3., 4., 3., 4., 3., 4.}, + {3, 2, 6}); + + auto tiled_1 = tile(data, {3}); + + auto data_2 = + array({4., 4., 9., 7., 7., 3., 6., 8., 8., 3., 8., 7.}, {3, 2, 2}); + + auto expected_2 = array( + {4., 4., 4., 4., 9., 7., 9., 7., 4., 4., 4., 4., 9., 7., 9., 7., 4., 4., + 4., 4., 9., 7., 9., 7., 7., 3., 7., 3., 6., 8., 6., 8., 7., 3., 7., 3., + 6., 8., 6., 8., 7., 3., 7., 3., 6., 8., 6., 8., 8., 3., 8., 3., 8., 7., + 8., 7., 8., 3., 8., 3., 8., 7., 8., 7., 8., 3., 8., 3., 8., 7., 8., 7.}, + {3, 6, 4}); + + auto tiled_2 = tile(data_2, {3, 2}); + + auto expected_3 = array( + {4., 4., 4., 4., 9., 7., 9., 7., 4., 4., 4., 4., 9., 7., 9., 7., 7., 3., + 7., 3., 6., 8., 6., 8., 7., 3., 7., 3., 6., 8., 6., 8., 8., 3., 8., 3., + 8., 7., 8., 7., 8., 3., 8., 3., 8., 7., 8., 7., 4., 4., 4., 4., 9., 7., + 9., 7., 4., 4., 4., 4., 9., 7., 9., 7., 7., 3., 7., 3., 6., 8., 6., 8., + 7., 3., 7., 3., 6., 8., 6., 8., 8., 3., 8., 3., 8., 7., 8., 7., 8., 3., + 8., 3., 8., 7., 8., 7., 4., 4., 4., 4., 9., 7., 9., 7., 4., 4., 4., 4., + 9., 7., 9., 7., 7., 3., 7., 3., 6., 8., 6., 8., 7., 3., 7., 3., 6., 8., + 6., 8., 8., 3., 8., 3., 8., 7., 8., 7., 8., 3., 8., 3., 8., 7., 8., 7.}, + {9, 4, 4}); + + auto tiled_3 = tile(data_2, {3, 2, 2}); + + auto expected_4 = array( + {4., 4., 4., 4., 9., 7., 9., 7., 4., 4., 4., 4., 9., 7., 9., 7., 7., 3., + 7., 3., 6., 8., 6., 8., 7., 3., 7., 3., 6., 8., 6., 8., 8., 3., 8., 3., + 8., 7., 8., 7., 8., 3., 8., 3., 8., 7., 8., 7., 4., 4., 4., 4., 9., 7., + 9., 7., 4., 4., 4., 4., 9., 7., 9., 7., 7., 3., 7., 3., 6., 8., 6., 8., + 7., 3., 7., 3., 6., 8., 6., 8., 8., 3., 8., 3., 8., 7., 8., 7., 8., 3., + 8., 3., 8., 7., 8., 7., 4., 4., 4., 4., 9., 7., 9., 7., 4., 4., 4., 4., + 9., 7., 9., 7., 7., 3., 7., 3., 6., 8., 6., 8., 7., 3., 7., 3., 6., 8., + 6., 8., 8., 3., 8., 3., 8., 7., 8., 7., 8., 3., 8., 3., 8., 7., 8., 7., + 4., 4., 4., 4., 9., 7., 9., 7., 4., 4., 4., 4., 9., 7., 9., 7., 7., 3., + 7., 3., 6., 8., 6., 8., 7., 3., 7., 3., 6., 8., 6., 8., 8., 3., 8., 3., + 8., 7., 8., 7., 8., 3., 8., 3., 8., 7., 8., 7., 4., 4., 4., 4., 9., 7., + 9., 7., 4., 4., 4., 4., 9., 7., 9., 7., 7., 3., 7., 3., 6., 8., 6., 8., + 7., 3., 7., 3., 6., 8., 6., 8., 8., 3., 8., 3., 8., 7., 8., 7., 8., 3., + 8., 3., 8., 7., 8., 7., 4., 4., 4., 4., 9., 7., 9., 7., 4., 4., 4., 4., + 9., 7., 9., 7., 7., 3., 7., 3., 6., 8., 6., 8., 7., 3., 7., 3., 6., 8., + 6., 8., 8., 3., 8., 3., 8., 7., 8., 7., 8., 3., 8., 3., 8., 7., 8., 7.}, + {3, 6, 4, 4}); + + auto tiled_4 = tile(data_2, {3, 2, 2, 2}); + + auto expected_5 = + array({4., 4., 9., 7., 7., 3., 6., 8., 8., 3., 8., 7.}, {3, 2, 2}); + auto tiled_5 = tile(data_2, {1, 1, 1}); + + auto expected_6 = + array({4., 4., 9., 7., 7., 3., 6., 8., 8., 3., 8., 7.}, {1, 3, 2, 2}); + auto tiled_6 = tile(data_2, {1, 1, 1, 1}); + + CHECK(array_equal(expected_1, tiled_1).item()); + CHECK(array_equal(expected_2, tiled_2).item()); + CHECK(array_equal({expected_3}, {tiled_3}).item()); + CHECK(array_equal({expected_4}, {tiled_4}).item()); + CHECK(array_equal({expected_5}, {tiled_5}).item()); + CHECK(array_equal({expected_6}, {tiled_6}).item()); + + CHECK_THROWS_AS(tile(data_2, {-1, 2}), std::invalid_argument); +} From 583d3734e9cbe7be3e89fc7ba7df482ce0e63eca Mon Sep 17 00:00:00 2001 From: Bahaatbb Date: Mon, 1 Jan 2024 00:10:08 +0300 Subject: [PATCH 2/3] Improved tile using broadcasting --- mlx/ops.cpp | 30 ++++++++++++++++-------------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 8e153e33ac..ee465fa33d 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -759,6 +759,13 @@ array tile(const array& a, std::vector repeats, StreamOrDevice s) { return a; } + for (size_t i = 0; i < repeats.size(); i++) { + if (repeats[i] < 0) { + throw std::invalid_argument( + "[tile] Negative repeat value found at index " + std::to_string(i)); + } + } + int d = repeats.size(); array arr = copy(a, s); @@ -768,24 +775,19 @@ array tile(const array& a, std::vector repeats, StreamOrDevice s) { arr = expand_dims(arr, std::vector(d - arr.ndim(), 0), s); } - std::vector shape_out; - for (size_t i = 0; i < arr.shape().size(); i++) { - shape_out.push_back(arr.shape()[i] * repeats[i]); - } - - int n = arr.size(); - std::vector old_shape = arr.shape(); - if (n > 0) { - for (size_t i = 0; i < old_shape.size(); i++) { - if (repeats[i] != 1) { - arr = repeat(reshape(arr, {-1, n}, s), repeats[i], 0, s); - } - n /= old_shape[i]; + if (repeats[i] != 1) { + arr = expand_dims(arr, i, s); + std::vector new_shape(arr.shape()); + new_shape[i] *= repeats[i]; + arr = broadcast_to(arr, new_shape, s); + new_shape[i] = -1; + new_shape.erase(new_shape.begin() + i + 1); + arr = reshape(arr, new_shape, s); } } - return reshape(arr, shape_out, s); + return arr; } /** Pad an array with a constant value */ From d9a6b54acd164537a826135ecb2ca2ce9a1b64e6 Mon Sep 17 00:00:00 2001 From: Bahaatbb Date: Wed, 3 Jan 2024 17:52:18 +0300 Subject: [PATCH 3/3] resolved some conflicts --- ACKNOWLEDGMENTS.md | 1 + 1 file changed, 1 insertion(+) diff --git a/ACKNOWLEDGMENTS.md b/ACKNOWLEDGMENTS.md index 15648e6e00..bec8b359b1 100644 --- a/ACKNOWLEDGMENTS.md +++ b/ACKNOWLEDGMENTS.md @@ -11,6 +11,7 @@ MLX was developed with contributions from the following individuals: - Juarez Bochi: Fixed bug in cross attention. - Justin Deschenaux: Sine, Cosine, arange, randint, truncated normal, bernoulli, lion optimizer, Dropout2d, linear and logistic regression python example. - Diogo Da Cruz: Added `tri`, `tril`, `triu`, `tensordot` and safetensor support +- Bahaa Eddin Tabbakha: Added `repeat`, `tile`. # Third-Party Software