Export: preserve Dtype state values in export callback arguments (#3145)

Co-authored-by: Awni Hannun <[email protected]>

python/src/export.cpp

@@ -164,6 +164,8 @@ void init_export(nb::module_& m) {
       nb::kw_only(),
       "shapeless"_a = false,
       "kwargs"_a,
+      nb::sig(
+          "def export_function(file_or_callback: Union[str, Callable], fun: Callable, *args, shapeless: bool = False, **kwargs) -> None"),
       R"pbdoc(
         Export an MLX function.
 
@@ -178,8 +180,8 @@ void init_export(nb::module_& m) {
           versions of MLX may not be compatible with future versions.
 
         Args:
-            file (str or Callable): Either a file path to export the function
-              to or a callback.
+            file_or_callback (str or Callable): Either a file path to export
+              the function to or a callback.
             fun (Callable): A function which takes as input zero or more
               :class:`array` and returns one or more :class:`array`.
             *args (array): Example array inputs to the function.

python/tests/test_export_import.py

@@ -501,13 +501,14 @@ class TestExportImport(mlx_tests.MLXTestCase):
     def test_export_with_callback(self):
 
         def fn(x, y):
-            return mx.log(mx.abs(x - y))
+            return mx.log(mx.abs(x - y)).astype(mx.int32)
 
         n_in = None
         n_out = None
         n_const = None
         keywords = None
         primitives = []
+        primitive_args = []
 
         def callback(args):
             nonlocal n_in, n_out, n_const, keywords, primitives
@@ -522,6 +523,7 @@ class TestExportImport(mlx_tests.MLXTestCase):
                 n_const = len(args["constants"])
             elif t == "primitive":
                 primitives.append(args["name"])
+                primitive_args.append(args["arguments"])
 
         mx.export_function(callback, fn, mx.array(1.0), y=mx.array(1.0))
         self.assertEqual(n_in, 2)
@@ -529,7 +531,11 @@ class TestExportImport(mlx_tests.MLXTestCase):
         self.assertEqual(n_const, 0)
         self.assertEqual(len(keywords), 1)
         self.assertEqual(keywords[0][0], "y")
-        self.assertEqual(primitives, ["Subtract", "Abs", "Log"])
+        self.assertEqual(primitives, ["Subtract", "Abs", "Log", "AsType"])
+        self.assertEqual(primitive_args[0], [])
+        self.assertEqual(primitive_args[1], [])
+        self.assertEqual(primitive_args[2], [2])
+        self.assertEqual(primitive_args[3], [mx.int32])
 
     @unittest.skipIf(not mx.is_available(mx.gpu), "No GPU available")
     def test_export_import_custom_kernel(self):