From eff0e31f004ef7f87a8d3abdddbbf5cee34deac0 Mon Sep 17 00:00:00 2001
From: Awni Hannun <awni@apple.com>
Date: Tue, 2 Dec 2025 11:24:40 -0800
Subject: [PATCH] Fix export scatters (#2852)

---
 mlx/export.cpp                     |  1 +
 mlx/primitives.h                   | 10 +++++-----
 python/tests/test_export_import.py | 13 +++++++++++++
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/mlx/export.cpp b/mlx/export.cpp
index 01655fca..9a773cd7 100644
--- a/mlx/export.cpp
+++ b/mlx/export.cpp
@@ -382,6 +382,7 @@ struct PrimitiveFactory {
       SERIALIZE_PRIMITIVE(LogicalOr),
       SERIALIZE_PRIMITIVE(LogAddExp),
       SERIALIZE_PRIMITIVE(LogSumExp),
+      SERIALIZE_PRIMITIVE(MaskedScatter),
       SERIALIZE_PRIMITIVE(Matmul),
       SERIALIZE_PRIMITIVE(Maximum),
       SERIALIZE_PRIMITIVE(Minimum),
diff --git a/mlx/primitives.h b/mlx/primitives.h
index 87aad33b..3d9f60b4 100644
--- a/mlx/primitives.h
+++ b/mlx/primitives.h
@@ -1871,13 +1871,13 @@ class Scatter : public UnaryPrimitive {
   const char* name() const override {
     switch (reduce_type_) {
       case Sum:
-        return "ScatterSum";
+        return "Scatter Sum";
       case Prod:
-        return "ScatterProd";
+        return "Scatter Prod";
       case Min:
-        return "ScatterMin";
+        return "Scatter Min";
       case Max:
-        return "ScatterMax";
+        return "Scatter Max";
       case None:
         return "Scatter";
     }
@@ -1910,7 +1910,7 @@ class ScatterAxis : public UnaryPrimitive {
   const char* name() const override {
     switch (reduce_type_) {
       case Sum:
-        return "ScatterAxisSum";
+        return "ScatterAxis Sum";
       case None:
         return "ScatterAxis";
     }
diff --git a/python/tests/test_export_import.py b/python/tests/test_export_import.py
index 6e280a79..c8f99bbb 100644
--- a/python/tests/test_export_import.py
+++ b/python/tests/test_export_import.py
@@ -596,6 +596,19 @@ class TestExportImport(mlx_tests.MLXTestCase):
         for y in ys:
             self.assertEqual(imported(y)[0].item(), fun(y).item())
 
+    def test_export_import_scatter_sum(self):
+        def fun(x, y, z):
+            return x.at[y].add(z)
+
+        x = mx.array([1, 2, 3])
+        y = mx.array([0, 0, 1])
+        z = mx.array([1, 1, 1])
+        path = os.path.join(self.test_dir, "fn.mlxfn")
+        mx.export_function(path, fun, x, y, z)
+
+        imported = mx.import_function(path)
+        self.assertTrue(mx.array_equal(imported(x, y, z)[0], fun(x, y, z)))
+
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()