zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-24 01:17:26 +08:00
Code Issues Actions 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
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 28 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mlx/export.cpp
View File

@ -266,6 +266,7 @@ struct PrimitiveFactory {
SERIALIZE_PRIMITIVE(Floor),
SERIALIZE_PRIMITIVE(Full),
SERIALIZE_PRIMITIVE(Gather),
SERIALIZE_PRIMITIVE(GatherAxis),
SERIALIZE_PRIMITIVE(GatherMM),
SERIALIZE_PRIMITIVE(Greater),
SERIALIZE_PRIMITIVE(GreaterEqual),
@ -307,6 +308,7 @@ struct PrimitiveFactory {
"CumMax",
"CumLogaddexp"),
SERIALIZE_PRIMITIVE(Scatter),
SERIALIZE_PRIMITIVE(ScatterAxis),
SERIALIZE_PRIMITIVE(Select),
SERIALIZE_PRIMITIVE(Sigmoid),
SERIALIZE_PRIMITIVE(Sign),

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()