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

Added mx.stack c++ frontend impl (#123)

Browse Source 
* stack C++ operation + python bindings
This commit is contained in:
Jason
2023-12-14 16:21:19 -05:00
committed by GitHub
parent e5851e52b1
commit e28b57e371
9 changed files with 191 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
python/tests/test_ops.py
View File

@@ -1371,6 +1371,37 @@ class TestOps(mlx_tests.MLXTestCase):
np_eye_matrix = np.eye(5, 6, k=-2)
self.assertTrue(np.array_equal(eye_matrix, np_eye_matrix))
def test_stack(self):
a = mx.ones((2,))
np_a = np.ones((2,))
b = mx.ones((2,))
np_b = np.ones((2,))
# One dimensional stack axis=0
c = mx.stack([a, b])
np_c = np.stack([np_a, np_b])
self.assertTrue(np.array_equal(c, np_c))
# One dimensional stack axis=1
c = mx.stack([a, b], axis=1)
np_c = np.stack([np_a, np_b], axis=1)
self.assertTrue(np.array_equal(c, np_c))
a = mx.ones((1, 2))
np_a = np.ones((1, 2))
b = mx.ones((1, 2))
np_b = np.ones((1, 2))
# Two dimensional stack axis=0
c = mx.stack([a, b])
np_c = np.stack([np_a, np_b])
self.assertTrue(np.array_equal(c, np_c))
# Two dimensional stack axis=1
c = mx.stack([a, b], axis=1)
np_c = np.stack([np_a, np_b], axis=1)
self.assertTrue(np.array_equal(c, np_c))
if __name__ == "__main__":
unittest.main()