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

fix build w/ flatten (#195)

Browse Source
This commit is contained in:
Awni Hannun
2023-12-17 11:58:45 -08:00
committed by GitHub
parent 52e1589a52
commit 90d04072b7
3 changed files with 30 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
tests/ops_tests.cpp
View File

@@ -73,6 +73,12 @@ TEST_CASE("test flatten") {
// Check start > end throws
CHECK_THROWS(flatten(x, 2, 1));
// Check start >= ndim throws
CHECK_THROWS(flatten(x, 5, 6));
// Check end < 0 throws
CHECK_THROWS(flatten(x, -5, -4));
// Check scalar flattens to 1D
x = array(1);
CHECK_EQ(flatten(x, -3, -1).shape(), std::vector<int>({1}));