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

implemented Flatten Module (#149)

Browse Source 
* implemented flatten op

---------

Co-authored-by: Awni Hannun <awni@apple.com>
This commit is contained in:
__mo_san__
2023-12-17 06:54:37 +01:00
committed by GitHub
parent eebd7c275d
commit 52e1589a52
8 changed files with 113 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
python/tests/test_ops.py
View File

@@ -1426,6 +1426,15 @@ class TestOps(mlx_tests.MLXTestCase):
np_c = np.stack([np_a, np_b], axis=1)
self.assertTrue(np.array_equal(c, np_c))
def test_flatten(self):
x = mx.zeros([2, 3, 4])
self.assertEqual(mx.flatten(x).shape, [2 * 3 * 4])
self.assertEqual(mx.flatten(x, start_axis=1).shape, [2, 3 * 4])
self.assertEqual(mx.flatten(x, end_axis=1).shape, [2 * 3, 4])
self.assertEqual(x.flatten().shape, [2 * 3 * 4])
self.assertEqual(x.flatten(start_axis=1).shape, [2, 3 * 4])
self.assertEqual(x.flatten(end_axis=1).shape, [2 * 3, 4])
if __name__ == "__main__":
unittest.main()