Floor and Ceil (#150)

* Implements Floor and Ceil Ops
This commit is contained in:
Luca Arnaboldi
2023-12-14 19:00:23 +01:00
committed by GitHub
parent 1e0c78b970
commit b93c4cf378
14 changed files with 250 additions and 4 deletions

View File

@@ -334,6 +334,22 @@ class TestOps(mlx_tests.MLXTestCase):
expected = [1, -5, 10]
self.assertListEqual(mx.maximum(x, y).tolist(), expected)
def test_floor(self):
x = mx.array([-22.03, 19.98, -27, 9, 0.0, -np.inf, np.inf])
expected = [-23, 19, -27, 9, 0, -np.inf, np.inf]
self.assertListEqual(mx.floor(x).tolist(), expected)
with self.assertRaises(ValueError):
mx.floor(mx.array([22 + 3j, 19 + 98j]))
def test_ceil(self):
x = mx.array([-22.03, 19.98, -27, 9, 0.0, -np.inf, np.inf])
expected = [-22, 20, -27, 9, 0, -np.inf, np.inf]
self.assertListEqual(mx.ceil(x).tolist(), expected)
with self.assertRaises(ValueError):
mx.floor(mx.array([22 + 3j, 19 + 98j]))
def test_transpose_noargs(self):
x = mx.array([[0, 1, 1], [1, 0, 0]])