zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-04 15:04:40 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add support for repeat (#278)

Browse Source 
* 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>
This commit is contained in:
Bahaa
2023-12-28 00:11:38 +03:00
committed by GitHub
parent 4417e37ede
commit ff2b58e299
6 changed files with 149 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
tests/ops_tests.cpp
View File

@@ -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);
}