Comparing python objects (such as list/tuple) with mlx.core.array (#920)

* add implicit conversion of list to array for equality constraint

* add tests for array equality

* add test for tuple and array equality

* return False if __eq__ arg is list or tuple

* write tests for equality

* update the rule of comparison for __ge__/__gt__/__lt__/__le__

* add a helper function for detecting mlx.core.array

* return true in case fo inequality

* debug minor issue regarding detecting mlx array

* add tests for inequality comparisons

* add name for contribution

* reformat files using pre-commit

* update tests for float

* update tests for inequality

* raise exception in case of invalid comparisons

* use isinstance instead of string comparison

* replace "is_convirtable_to_array" with previous logic

* remove throwing exceptions for other operations

* just a comment

* minor changes for efficiency

* optimize a utils function

* change the function name

* Update ACKNOWLEDGMENTS.md

---------

Co-authored-by: Awni Hannun <awni.hannun@gmail.com>
This commit is contained in:
AmirHossein_Razlighi
2024-03-29 17:22:30 +03:30
committed by GitHub
parent 913b19329c
commit f48bc496c7
4 changed files with 143 additions and 8 deletions

View File

@@ -89,6 +89,121 @@ class TestDtypes(mlx_tests.MLXTestCase):
self.assertListEqual(list(z.shape), list(y.shape))
class TestEquality(mlx_tests.MLXTestCase):
def test_array_eq_array(self):
a = mx.array([1, 2, 3])
b = mx.array([1, 2, 3])
c = mx.array([1, 2, 4])
self.assertTrue(mx.all(a == b))
self.assertFalse(mx.all(a == c))
def test_array_eq_scalar(self):
a = mx.array([1, 2, 3])
b = 1
c = 4
d = 2.5
e = mx.array([1, 2.5, 3.25])
self.assertTrue(mx.any(a == b))
self.assertFalse(mx.all(a == c))
self.assertFalse(mx.all(a == d))
self.assertTrue(mx.any(a == e))
def test_list_equals_array(self):
a = mx.array([1, 2, 3])
b = [1, 2, 3]
c = [1, 2, 4]
# mlx array equality returns false if is compared with any kind of
# object which is not an mlx array
self.assertFalse(a == b)
self.assertFalse(a == c)
def test_tuple_equals_array(self):
a = mx.array([1, 2, 3])
b = (1, 2, 3)
c = (1, 2, 4)
# mlx array equality returns false if is compared with any kind of
# object which is not an mlx array
self.assertFalse(a == b)
self.assertFalse(a == c)
class TestInequality(mlx_tests.MLXTestCase):
def test_array_ne_array(self):
a = mx.array([1, 2, 3])
b = mx.array([1, 2, 3])
c = mx.array([1, 2, 4])
self.assertFalse(mx.any(a != b))
self.assertTrue(mx.any(a != c))
def test_array_ne_scalar(self):
a = mx.array([1, 2, 3])
b = 1
c = 4
d = 1.5
e = 2.5
f = mx.array([1, 2.5, 3.25])
self.assertFalse(mx.all(a != b))
self.assertTrue(mx.any(a != c))
self.assertTrue(mx.any(a != d))
self.assertTrue(mx.any(a != e))
self.assertFalse(mx.all(a != f))
def test_list_not_equals_array(self):
a = mx.array([1, 2, 3])
b = [1, 2, 3]
c = [1, 2, 4]
# mlx array inequality returns true if is compared with any kind of
# object which is not an mlx array
self.assertTrue(a != b)
self.assertTrue(a != c)
def test_tuple_not_equals_array(self):
a = mx.array([1, 2, 3])
b = (1, 2, 3)
c = (1, 2, 4)
# mlx array inequality returns true if is compared with any kind of
# object which is not an mlx array
self.assertTrue(a != b)
self.assertTrue(a != c)
def test_obj_inequality_array(self):
str_ = "hello"
a = mx.array([1, 2, 3])
lst_ = [1, 2, 3]
tpl_ = (1, 2, 3)
# check if object comparison(</>/<=/>=) with mlx array should throw an exception
# if not, the tests will fail
with self.assertRaises(ValueError):
a < str_
with self.assertRaises(ValueError):
a > str_
with self.assertRaises(ValueError):
a <= str_
with self.assertRaises(ValueError):
a >= str_
with self.assertRaises(ValueError):
a < lst_
with self.assertRaises(ValueError):
a > lst_
with self.assertRaises(ValueError):
a <= lst_
with self.assertRaises(ValueError):
a >= lst_
with self.assertRaises(ValueError):
a < tpl_
with self.assertRaises(ValueError):
a > tpl_
with self.assertRaises(ValueError):
a <= tpl_
with self.assertRaises(ValueError):
a >= tpl_
class TestArray(mlx_tests.MLXTestCase):
def test_array_basics(self):
x = mx.array(1)