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

Return empty array when repeats is 0 in mx.repeat (#681)

Browse Source 
* Return empty array when repeats is 0

* Add test case for repeats = 0
This commit is contained in:
Noah Farr
2024-02-14 02:49:31 +01:00
committed by GitHub
parent 2fdc2462c3
commit 0c65517e91
2 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/tests/test_ops.py
View File

@@ -1699,6 +1699,8 @@ class TestOps(mlx_tests.MLXTestCase):
def test_repeat(self):
# Setup data for the tests
data = mx.array([[[13, 3], [16, 6]], [[14, 4], [15, 5]], [[11, 1], [12, 2]]])
# Test repeat 0 times
self.assertCmpNumpy([data, 0], mx.repeat, np.repeat)
# Test repeat along axis 0
self.assertCmpNumpy([data, 2], mx.repeat, np.repeat, axis=0)
# Test repeat along axis 1