From f2f2d1645156938a982cbeaff87a2d83795905c0 Mon Sep 17 00:00:00 2001 From: Alex Skryl Date: Thu, 19 Feb 2026 10:07:28 -0600 Subject: [PATCH] Export: preserve Dtype state values in export callback arguments (#3145) Co-authored-by: Awni Hannun --- docs/src/usage/export.rst | 28 ++++++++++++++++++++++++++++ mlx/export.cpp | 2 ++ python/src/export.cpp | 6 ++++-- python/tests/test_export_import.py | 10 ++++++++-- 4 files changed, 42 insertions(+), 4 deletions(-) diff --git a/docs/src/usage/export.rst b/docs/src/usage/export.rst index e908b85f..ac1ec218 100644 --- a/docs/src/usage/export.rst +++ b/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 ----------------- diff --git a/mlx/export.cpp b/mlx/export.cpp index 85fc2774..1160e8cf 100644 --- a/mlx/export.cpp +++ b/mlx/export.cpp @@ -279,6 +279,8 @@ void extract_state(const T state, std::vector& unpacked_state) { unpacked_state.push_back(state); } else if constexpr (std::is_enum_v) { unpacked_state.push_back(static_cast(state)); + } else if constexpr (std::is_same_v) { + unpacked_state.push_back(state); } else if constexpr (is_iterable) { unpacked_state.push_back(state); } else if constexpr (is_pair || is_tuple) { diff --git a/python/src/export.cpp b/python/src/export.cpp index 3d3ffdf6..51c8a896 100644 --- a/python/src/export.cpp +++ b/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. diff --git a/python/tests/test_export_import.py b/python/tests/test_export_import.py index c8f99bbb..6f2b2ca4 100644 --- a/python/tests/test_export_import.py +++ b/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):