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

Allow None input to compiled functions (#2621)

Browse Source 
* Allow None input to compiled functions

* Allow None input to compiled functions
This commit is contained in:
Awni Hannun
2025-09-25 08:42:23 -07:00
committed by GitHub
parent 7c7e48dbd1
commit d4f4ff3c5e
2 changed files with 18 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
python/tests/test_compile.py
View File

@@ -1037,6 +1037,20 @@ class TestCompile(mlx_tests.MLXTestCase):
self.assertEqual(inner.__doc__, c_inner.__doc__)
self.assertEqual(inspect.signature(inner), inspect.signature(c_inner))
def test_compile_with_none(self):
@mx.compile
def fun(x, y):
if y is None:
return mx.abs(x - 2.0)
else:
return mx.abs(x + y)
out = fun(mx.array(1.0), None)
self.assertEqual(out.item(), 1.0)
out = fun(mx.array(1.0), mx.array(2.0))
self.assertEqual(out.item(), 3.0)
if __name__ == "__main__":
mlx_tests.MLXTestRunner()