zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-17 09:18:12 +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
2 changed files with 28 additions and 0 deletions
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),