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

fix export to work with gather/scatter axis (#2263)

Browse Source
This commit is contained in:
Awni Hannun
2025-06-09 20:37:27 -07:00
committed by GitHub
parent f8bad60609
commit 9ce77798b1
2 changed files with 28 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
python/tests/test_export_import.py
View File

@@ -286,6 +286,32 @@ class TestExportImport(mlx_tests.MLXTestCase):
with self.assertRaises(ValueError):
f2(mx.array(10), mx.array([5, 10, 20]))
def test_export_scatter_gather(self):
path = os.path.join(self.test_dir, "fn.mlxfn")
def fun(a, b):
return mx.take_along_axis(a, b, axis=0)
x = mx.random.uniform(shape=(4, 4))
y = mx.array([[0, 1, 2, 3], [1, 2, 0, 3]])
mx.export_function(path, fun, (x, y))
imported_fun = mx.import_function(path)
expected = fun(x, y)
out = imported_fun(x, y)[0]
self.assertTrue(mx.array_equal(expected, out))
def fun(a, b, c):
return mx.put_along_axis(a, b, c, axis=0)
x = mx.random.uniform(shape=(4, 4))
y = mx.array([[0, 1, 2, 3], [1, 2, 0, 3]])
z = mx.random.uniform(shape=(2, 4))
mx.export_function(path, fun, (x, y, z))
imported_fun = mx.import_function(path)
expected = fun(x, y, z)
out = imported_fun(x, y, z)[0]
self.assertTrue(mx.array_equal(expected, out))
if __name__ == "__main__":
unittest.main()