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

Add types, fix kwarg ordering bug + test

Browse Source
This commit is contained in:
Awni Hannun
2025-09-24 15:50:53 -07:00
parent 9fcfcf04c6
commit ad0dd9b5ba
4 changed files with 50 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
python/tests/test_export_import.py
View File

@@ -485,6 +485,19 @@ class TestExportImport(mlx_tests.MLXTestCase):
mx.array_equal(imported_fn(input_data)[0], model(input_data))
)
def test_export_kwarg_ordering(self):
path = os.path.join(self.test_dir, "fun.mlxfn")
def fn(x, y):
return x - y
mx.export_function(path, fn, x=mx.array(1.0), y=mx.array(1.0))
imported = mx.import_function(path)
out = imported(x=mx.array(2.0), y=mx.array(3.0))[0]
self.assertEqual(out.item(), -1.0)
out = imported(y=mx.array(2.0), x=mx.array(3.0))[0]
self.assertEqual(out.item(), 1.0)
if __name__ == "__main__":
mlx_tests.MLXTestRunner()