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

Added linspace (#181)

Browse Source 
* linspace ops support

---------

Co-authored-by: Awni Hannun <awni@apple.com>
This commit is contained in:
Abe Leininger
2023-12-18 19:57:55 -08:00
committed by GitHub
parent f4f6e17d45
commit e6872a4149
6 changed files with 92 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
python/tests/test_ops.py
View File

@@ -1491,6 +1491,27 @@ class TestOps(mlx_tests.MLXTestCase):
clipped = mx.clip(mx.array(a), mx.array(mins), mx.array(maxs))
self.assertTrue(np.array_equal(clipped, expected))
def test_linspace(self):
# Test default num = 50
a = mx.linspace(0, 1)
expected = mx.array(np.linspace(0, 1))
self.assertEqualArray(a, expected)
# Test int32 dtype
b = mx.linspace(0, 10, 5, mx.int64)
expected = mx.array(np.linspace(0, 10, 5, dtype=int))
self.assertEqualArray(b, expected)
# Test negative sequence with float start and stop
c = mx.linspace(-2.7, -0.7, 7)
expected = mx.array(np.linspace(-2.7, -0.7, 7))
self.assertEqualArray(c, expected)
# Test irrational step size of 1/9
d = mx.linspace(0, 1, 10)
expected = mx.array(np.linspace(0, 1, 10))
self.assertEqualArray(d, expected)
if __name__ == "__main__":
unittest.main()