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

fix clip (#1415)

Browse Source
This commit is contained in:
Awni Hannun
2024-09-14 16:09:09 -07:00
committed by GitHub
parent b3f52c9fbe
commit d6492b0163
3 changed files with 19 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
python/tests/test_ops.py
View File

@@ -2008,6 +2008,22 @@ class TestOps(mlx_tests.MLXTestCase):
clipped = mx.clip(mx.array(a), mx.array(mins), mx.array(maxs))
self.assertTrue(np.array_equal(clipped, expected))
# Check clip output types
a = mx.array([1, 2, 3], mx.int16)
out_t = mx.clip(a, a_min=0, a_max=5).dtype
self.assertEqual(out_t, mx.int16)
out_t = mx.clip(a, a_min=0.0, a_max=5).dtype
self.assertEqual(out_t, mx.float32)
a = mx.array([1, 2, 3], mx.float16)
out_t = mx.clip(a, a_min=0.0, a_max=5).dtype
self.assertEqual(out_t, mx.float16)
a = mx.array([1, 2, 3], mx.float16)
out_t = mx.clip(a, a_min=0.0, a_max=mx.array(1.0)).dtype
self.assertEqual(out_t, mx.float32)
def test_linspace(self):
# Test default num = 50
a = mx.linspace(0, 1)