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 diff --git a/docs/src/python/ops.rst b/docs/src/python/ops.rst index 4e399524e0..f167a917d9 100644 --- a/docs/src/python/ops.rst +++ b/docs/src/python/ops.rst @@ -104,6 +104,7 @@ Operations take_along_axis tan tanh + tile tensordot transpose tri diff --git a/mlx/ops.cpp b/mlx/ops.cpp index e1f593abaf..b76770073a 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -753,6 +753,43 @@ 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; + } + + 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); + + 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); + } + + for (size_t i = 0; i < arr.shape().size(); 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 arr; +} + /** Pad an array with a constant value */ array pad( const array& a, diff --git a/mlx/ops.h b/mlx/ops.h index 6516ad008d..21f5be9df5 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 c152eeb97c..1ce0eee408 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 65de09634b..901bdcff0c 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,50 @@ 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)) + + def test_tensordot(self): x = mx.arange(60.0).reshape(3, 4, 5) y = mx.arange(24.0).reshape(4, 3, 2) diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index ad767f5255..3dd6de9d2c 100644 --- a/tests/ops_tests.cpp +++ b/tests/ops_tests.cpp @@ -2279,6 +2279,80 @@ TEST_CASE("test repeat") { CHECK_THROWS_AS(repeat(data_3, -3, 0), std::invalid_argument); } +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); +} TEST_CASE("tensordot") { auto x = reshape(arange(60.), {3, 4, 5}); auto y = reshape(arange(24.), {4, 3, 2}); @@ -2314,4 +2388,4 @@ TEST_CASE("tensordot") { 65370.}, {3, 6}); CHECK(array_equal(z, expected).item()); -} \ No newline at end of file +}