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

fix shapeless export (#2148)

Browse Source
This commit is contained in:
Awni Hannun 2025-05-01 15:02:02 -07:00 committed by GitHub
parent a3a632d567
commit 9daa6b003f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 22 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
mlx/export.cpp
View File

@ -470,6 +470,9 @@ bool FunctionTable::match(
if (x.dtype() != y.dtype()) {
return false;
}
if (x.ndim() != y.ndim()) {
return false;
}
if (!shapeless && x.shape() != y.shape()) {
return false;
}

19
python/tests/test_export_import.py
View File

@ -242,6 +242,7 @@ class TestExportImport(mlx_tests.MLXTestCase):
def test_leaks(self):
path = os.path.join(self.test_dir, "fn.mlxfn")
mx.synchronize()
if mx.metal.is_available():
mem_pre = mx.get_active_memory()
else:
@ -267,6 +268,24 @@ class TestExportImport(mlx_tests.MLXTestCase):
self.assertEqual(mem_pre, mem_post)
def test_export_import_shapeless(self):
path = os.path.join(self.test_dir, "fn.mlxfn")
def fun(*args):
return sum(args)
with mx.exporter(path, fun, shapeless=True) as exporter:
exporter(mx.array(1))
exporter(mx.array(1), mx.array(2))
exporter(mx.array(1), mx.array(2), mx.array(3))
f2 = mx.import_function(path)
self.assertEqual(f2(mx.array(1))[0].item(), 1)
self.assertEqual(f2(mx.array(1), mx.array(1))[0].item(), 2)
self.assertEqual(f2(mx.array(1), mx.array(1), mx.array(1))[0].item(), 3)
with self.assertRaises(ValueError):
f2(mx.array(10), mx.array([5, 10, 20]))
if __name__ == "__main__":
unittest.main()