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

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

docs/src/usage/export.rst

@@ -155,6 +155,34 @@ parameters, pass them as inputs to the ``call`` wrapper:
    mx.export_function("model.mlxfn", call, (mx.zeros(4),), params)
 
 
+Exporting with a Callback
+-------------------------
+
+To inspect the exported graph, you can pass a callback instead of a file path
+to :func:`export_function`.
+
+.. code-block:: python
+
+  def fun(x):
+    return x.astype(mx.int32)
+
+  def callback(args):
+    print(args)
+
+  mx.export_function(callback, fun, mx.array([1.0, 2.0]))
+
+The argument to the callback (``args``) is a dictionary which includes a
+``type`` field. The possible types are:
+
+* ``"inputs"``: The ordered positional inputs to the exported function
+* ``"keyword_inputs"``: The keyword specified inputs to the exported function
+* ``"outputs"``: The ordered outputs of the exported function
+* ``"constants"``: Any graph constants
+* ``"primitives"``: Inner graph nodes representating the operations
+
+Each type has additional fields in the ``args`` dictionary.
+
+
 Shapeless Exports
 -----------------
 

mlx/export.cpp

@@ -279,6 +279,8 @@ void extract_state(const T state, std::vector<StateT>& unpacked_state) {
     unpacked_state.push_back(state);
   } else if constexpr (std::is_enum_v<T>) {
     unpacked_state.push_back(static_cast<int>(state));
+  } else if constexpr (std::is_same_v<T, Dtype>) {
+    unpacked_state.push_back(state);
   } else if constexpr (is_iterable<T>) {
     unpacked_state.push_back(state);
   } else if constexpr (is_pair<T> || is_tuple<T>) {

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