L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

implemented Flatten Module (#149)

Browse Source 
* implemented flatten op

---------

Co-authored-by: Awni Hannun <[email protected]>
This commit is contained in:
__mo_san__
2023-12-16 21:54:37 -08:00
committed by GitHub
co-authored by Awni Hannun
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
python/tests/test_ops.py
+9
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()