diff --git a/docs/src/usage/export.rst b/docs/src/usage/export.rst index ac1ec218bb..89176344db 100644 --- a/docs/src/usage/export.rst +++ b/docs/src/usage/export.rst @@ -109,6 +109,33 @@ keyword arguments when calling the imported function. out, = imported_fun(x, z=y) +Saving Metadata +--------------- + +You can save metadata, such as a model configuration, alongside an exported +function. The metadata is a dictionary of string keys with string, integer, or +float values: + +.. code-block:: python + + def fun(x, y): + return x + y + + x = mx.array(1.0) + y = mx.array(1.0) + metadata = {"description": "adds two arrays", "version": 1, "scale": 0.5} + mx.export_function("add.mlxfn", fun, x, y, metadata=metadata) + +Pass ``return_metadata=True`` to read the metadata back when importing: + +.. code-block:: python + + imported_fun, metadata = mx.import_function("add.mlxfn", return_metadata=True) + + # Prints: adds two arrays + print(metadata["description"]) + + Exporting Modules ----------------- diff --git a/mlx/export.cpp b/mlx/export.cpp index 4777fd3048..710165a8f2 100644 --- a/mlx/export.cpp +++ b/mlx/export.cpp @@ -260,6 +260,55 @@ array deserialize(Reader& is) { return array(std::move(shape), type, nullptr, std::vector{}); } +enum class MetadataValueType { String = 0, Int = 1, Float = 2 }; + +void serialize( + Writer& os, + const std::unordered_map& metadata) { + serialize(os, static_cast(metadata.size())); + for (auto& [key, value] : metadata) { + serialize(os, key); + std::visit( + [&os](auto&& v) { + using T = std::decay_t; + if constexpr (std::is_same_v) { + serialize(os, MetadataValueType::String); + } else if constexpr (std::is_same_v) { + serialize(os, MetadataValueType::Int); + } else if constexpr (std::is_same_v) { + serialize(os, MetadataValueType::Float); + } + serialize(os, v); + }, + value); + } +} + +template <> +std::unordered_map deserialize(Reader& is) { + std::unordered_map metadata; + auto size = deserialize(is); + metadata.reserve(size); + for (uint64_t i = 0; i < size; ++i) { + auto key = deserialize(is); + switch (deserialize(is)) { + case MetadataValueType::String: + metadata.emplace(std::move(key), deserialize(is)); + break; + case MetadataValueType::Int: + metadata.emplace(std::move(key), deserialize(is)); + break; + case MetadataValueType::Float: + metadata.emplace(std::move(key), deserialize(is)); + break; + default: + throw std::runtime_error( + "[import_function] Unknown metadata value type."); + } + } + return metadata; +} + template constexpr bool has_state = false; @@ -521,10 +570,15 @@ struct PrimitiveFactory { }; }; -void write_header(Writer& os, int count, bool shapeless) { +void write_header( + Writer& os, + int count, + bool shapeless, + const std::unordered_map& metadata) { serialize(os, std::string(version())); serialize(os, count); serialize(os, shapeless); + serialize(os, metadata); } // A struct to hold and retrieve the graphs that are exported / imported @@ -674,14 +728,16 @@ FunctionTable::Function* FunctionTable::find( FunctionExporter::FunctionExporter( const std::string& file, std::function(const Args&, const Kwargs&)> fun, - bool shapeless) + bool shapeless, + std::unordered_map metadata) : os(file), fun(std::move(fun)), - ftable(std::make_shared(shapeless)) { + ftable(std::make_shared(shapeless)), + metadata_(std::move(metadata)) { if (!os.is_open()) { throw std::runtime_error("[export_function] Failed to open " + file); } - write_header(os, count, shapeless); + write_header(os, count, shapeless, metadata_); } FunctionExporter::FunctionExporter( @@ -819,7 +875,7 @@ void FunctionExporter::export_function(const Args& args, const Kwargs& kwargs) { // Update the header auto pos = os.tell(); os.seek(0); - write_header(os, count, ftable->shapeless); + write_header(os, count, ftable->shapeless, metadata_); os.seek(pos); serialize(os, kwarg_keys); @@ -898,44 +954,51 @@ void FunctionExporter::operator()(const Args& args, const Kwargs& kwargs) { FunctionExporter exporter( const std::string& file, const std::function(const Args&)>& fun, - bool shapeless /* = false */) { + bool shapeless /* = false */, + const std::unordered_map& metadata /* = {} */) { return FunctionExporter{ file, [fun](const Args& args, const Kwargs&) { return fun(args); }, - shapeless}; + shapeless, + metadata}; } FunctionExporter exporter( const std::string& file, const std::function(const Kwargs&)>& fun, - bool shapeless /* = false */) { + bool shapeless /* = false */, + const std::unordered_map& metadata /* = {} */) { return exporter( file, [fun](const Args&, const Kwargs kwargs) { return fun(kwargs); }, - shapeless); + shapeless, + metadata); } FunctionExporter exporter( const std::string& file, const std::function(const Args&, const Kwargs&)>& fun, - bool shapeless /* = false */) { - return FunctionExporter{file, fun, shapeless}; + bool shapeless /* = false */, + const std::unordered_map& metadata /* = {} */) { + return FunctionExporter{file, fun, shapeless, metadata}; } void export_function( const std::string& file, const std::function(const Args&)>& fun, const Args& args, - bool shapeless /* = false */) { - exporter(file, fun, shapeless)(args); + bool shapeless /* = false */, + const std::unordered_map& metadata /* = {} */) { + exporter(file, fun, shapeless, metadata)(args); } void export_function( const std::string& file, const std::function(const Kwargs&)>& fun, const Kwargs& kwargs, - bool shapeless /* = false */) { - exporter(file, fun, shapeless)(kwargs); + bool shapeless /* = false */, + const std::unordered_map& metadata /* = {} */) { + exporter(file, fun, shapeless, metadata)(kwargs); } void export_function( @@ -943,8 +1006,9 @@ void export_function( const std::function(const Args&, const Kwargs&)>& fun, const Args& args, const Kwargs& kwargs, - bool shapeless /* = false */) { - exporter(file, fun, shapeless)(args, kwargs); + bool shapeless /* = false */, + const std::unordered_map& metadata /* = {} */) { + exporter(file, fun, shapeless, metadata)(args, kwargs); } FunctionExporter exporter( @@ -1054,6 +1118,7 @@ ImportedFunction::ImportedFunction(const std::string& file) auto mlx_version = deserialize(is); auto function_count = deserialize(is); ftable->shapeless = deserialize(is); + metadata_ = deserialize>(is); std::unordered_map constants; auto import_one = [&]() { diff --git a/mlx/export.h b/mlx/export.h index 5532f7c818..cb7252a03e 100644 --- a/mlx/export.h +++ b/mlx/export.h @@ -14,6 +14,11 @@ namespace mlx::core { using Args = std::vector; using Kwargs = std::unordered_map; +// Metadata values saved alongside an exported function. The integer and +// floating point alternatives are 64-bit to hold Python int and float without +// narrowing (e.g. a parameter count above 2^31, or a non-power-of-two float). +using MetadataValue = std::variant; + // Possible types for a Primitive's state using StateT = std::variant< bool, @@ -50,17 +55,20 @@ struct FunctionExporter; MLX_API FunctionExporter exporter( const std::string& file, const std::function(const Args&)>& fun, - bool shapeless = false); + bool shapeless = false, + const std::unordered_map& metadata = {}); MLX_API FunctionExporter exporter( const std::string& file, const std::function(const Kwargs&)>& fun, - bool shapeless = false); + bool shapeless = false, + const std::unordered_map& metadata = {}); MLX_API FunctionExporter exporter( const std::string& path, const std::function(const Args&, const Kwargs&)>& fun, - bool shapeless = false); + bool shapeless = false, + const std::unordered_map& metadata = {}); /** * Export a function to a file. @@ -69,20 +77,23 @@ MLX_API void export_function( const std::string& file, const std::function(const Args&)>& fun, const Args& args, - bool shapeless = false); + bool shapeless = false, + const std::unordered_map& metadata = {}); MLX_API void export_function( const std::string& file, const std::function(const Kwargs&)>& fun, const Kwargs& kwargs, - bool shapeless = false); + bool shapeless = false, + const std::unordered_map& metadata = {}); MLX_API void export_function( const std::string& file, const std::function(const Args&, const Kwargs&)>& fun, const Args& args, const Kwargs& kwargs, - bool shapeless = false); + bool shapeless = false, + const std::unordered_map& metadata = {}); struct ImportedFunction; diff --git a/mlx/export_impl.h b/mlx/export_impl.h index 467a5f0d6c..f8cfa3a327 100644 --- a/mlx/export_impl.h +++ b/mlx/export_impl.h @@ -27,17 +27,20 @@ struct MLX_API FunctionExporter { friend MLX_API FunctionExporter exporter( const std::string&, const std::function(const Args&)>&, - bool shapeless); + bool shapeless, + const std::unordered_map& metadata); friend MLX_API FunctionExporter exporter( const std::string&, const std::function(const Kwargs&)>&, - bool shapeless); + bool shapeless, + const std::unordered_map& metadata); friend MLX_API FunctionExporter exporter( const std::string&, const std::function(const Args&, const Kwargs&)>&, - bool shapeless); + bool shapeless, + const std::unordered_map& metadata); friend MLX_API FunctionExporter exporter( const ExportCallback&, @@ -57,7 +60,8 @@ struct MLX_API FunctionExporter { FunctionExporter( const std::string& file, std::function(const Args&, const Kwargs&)> fun, - bool shapeless); + bool shapeless, + std::unordered_map metadata); FunctionExporter( const ExportCallback& callback, @@ -77,6 +81,7 @@ struct MLX_API FunctionExporter { int count{0}; bool closed{false}; std::shared_ptr ftable; + std::unordered_map metadata_; }; struct MLX_API ImportedFunction { @@ -88,12 +93,18 @@ struct MLX_API ImportedFunction { std::vector operator()(const Kwargs& kwargs) const; std::vector operator()(const Args& args, const Kwargs& kwargs) const; + // The metadata saved with the function when it was exported. + const std::unordered_map& metadata() const { + return metadata_; + } + private: ImportedFunction(const std::string& file); friend MLX_API ImportedFunction import_function(const std::string&); ImportedFunction(); std::shared_ptr ftable; + std::unordered_map metadata_; }; } // namespace mlx::core diff --git a/python/src/export.cpp b/python/src/export.cpp index 51c8a8965c..1948765b05 100644 --- a/python/src/export.cpp +++ b/python/src/export.cpp @@ -138,6 +138,8 @@ void init_export(nb::module_& m) { const nb::callable& fun, const nb::args& args, bool shapeless, + const std::optional< + std::unordered_map>& metadata, const nb::kwargs& kwargs) { auto [args_, kwargs_] = validate_and_extract_inputs(args, kwargs, "[export_function]"); @@ -147,8 +149,15 @@ void init_export(nb::module_& m) { wrap_export_function(fun), args_, kwargs_, - shapeless); + shapeless, + metadata.value_or( + std::unordered_map{})); } else { + if (metadata && !metadata->empty()) { + throw std::invalid_argument( + "[export_function] The metadata argument is only supported " + "when exporting to a file, not when using a callback."); + } auto callback = nb::cast(file_or_callback); auto wrapped_callback = [callback](const mx::ExportCallbackInput& input) { @@ -163,9 +172,10 @@ void init_export(nb::module_& m) { "args"_a, nb::kw_only(), "shapeless"_a = false, + "metadata"_a = nb::none(), "kwargs"_a, nb::sig( - "def export_function(file_or_callback: Union[str, Callable], fun: Callable, *args, shapeless: bool = False, **kwargs) -> None"), + "def export_function(file_or_callback: Union[str, Callable], fun: Callable, *args, shapeless: bool = False, metadata: Optional[dict[str, Union[str, int, float]]] = None, **kwargs) -> None"), R"pbdoc( Export an MLX function. @@ -187,6 +197,10 @@ void init_export(nb::module_& m) { *args (array): Example array inputs to the function. shapeless (bool, optional): Whether or not the function allows inputs with variable shapes. Default: ``False``. + metadata (dict, optional): A dictionary of string keys and string, + integer, or float values to save alongside the function. Only + supported when exporting to a file. Read it back with + :func:`import_function`. Default: ``None``. **kwargs (array): Additional example keyword array inputs to the function. @@ -203,17 +217,25 @@ void init_export(nb::module_& m) { )pbdoc"); m.def( "import_function", - [](const std::string& file) { - return nb::cpp_function( - [fn = mx::import_function(file)]( + [](const std::string& file, bool return_metadata) -> nb::object { + auto imported = mx::import_function(file); + auto metadata = imported.metadata(); + auto fn = nb::cpp_function( + [imported = std::move(imported)]( const nb::args& args, const nb::kwargs& kwargs) { auto [args_, kwargs_] = validate_and_extract_inputs( args, kwargs, "[import_function::call]"); - return nb::tuple(nb::cast(fn(args_, kwargs_))); + return nb::tuple(nb::cast(imported(args_, kwargs_))); }); + if (return_metadata) { + return nb::make_tuple(fn, nb::cast(metadata)); + } + return fn; }, "file"_a, - nb::sig("def import_function(file: str) -> Callable"), + "return_metadata"_a = false, + nb::sig( + "def import_function(file: str, return_metadata: bool = False) -> Union[Callable, tuple[Callable, dict[str, Union[str, int, float]]]]"), R"pbdoc( Import a function from a file. @@ -230,9 +252,14 @@ void init_export(nb::module_& m) { Args: file (str): The file path to import the function from. + return_metadata (bool, optional): If ``True`` also return the + metadata saved with the function as a dictionary. Default: + ``False``. Returns: - Callable: The imported function. + Callable: The imported function. If ``return_metadata`` is ``True`` + a tuple of the imported function and the metadata dictionary is + returned instead. Example: >>> fn = mx.import_function("function.mlxfn") @@ -274,14 +301,25 @@ void init_export(nb::module_& m) { m.def( "exporter", - [](const std::string& file, nb::callable fun, bool shapeless) { + [](const std::string& file, + nb::callable fun, + bool shapeless, + const std::optional< + std::unordered_map>& metadata) { return PyFunctionExporter{ - mx::exporter(file, wrap_export_function(fun), shapeless), fun}; + mx::exporter( + file, + wrap_export_function(fun), + shapeless, + metadata.value_or( + std::unordered_map{})), + fun}; }, "file"_a, "fun"_a, nb::kw_only(), "shapeless"_a = false, + "metadata"_a = nb::none(), R"pbdoc( Make a callable object to export multiple traces of a function to a file. @@ -295,6 +333,9 @@ void init_export(nb::module_& m) { file (str): File path to export the function to. shapeless (bool, optional): Whether or not the function allows inputs with variable shapes. Default: ``False``. + metadata (dict, optional): A dictionary of string keys and string, + integer, or float values to save alongside the function. Read it + back with :func:`import_function`. Default: ``None``. Example: diff --git a/python/tests/test_export_import.py b/python/tests/test_export_import.py index c45dd4a606..77ca9adf42 100644 --- a/python/tests/test_export_import.py +++ b/python/tests/test_export_import.py @@ -704,6 +704,50 @@ def fun(x, y, z): imported = mx.import_function(path) self.assertTrue(mx.array_equal(imported(x, y, z)[0], fun(x, y, z))) + def test_export_import_metadata(self): + path = os.path.join(self.test_dir, "fn.mlxfn") + + def fun(x): + return mx.abs(x) + + x = mx.array([1.0, -2.0, 3.0]) + # A value above 2**31 and a non-exact float, to catch narrowing + metadata = {"name": "model", "params": 7_000_000_000, "lr": 0.1} + + mx.export_function(path, fun, x, metadata=metadata) + + # No metadata returned by default, callable still works + imported = mx.import_function(path) + self.assertTrue(mx.array_equal(imported(x)[0], fun(x))) + + # Metadata read back with value types and exact values preserved + imported, imported_metadata = mx.import_function(path, return_metadata=True) + self.assertEqual(imported_metadata, metadata) + self.assertIsInstance(imported_metadata["name"], str) + self.assertIsInstance(imported_metadata["params"], int) + self.assertIsInstance(imported_metadata["lr"], float) + self.assertTrue(mx.array_equal(imported(x)[0], fun(x))) + + # No metadata gives an empty dict + mx.export_function(path, fun, x) + _, imported_metadata = mx.import_function(path, return_metadata=True) + self.assertEqual(imported_metadata, {}) + + # Metadata survives the per-trace header rewrite of a multi-trace export + with mx.exporter(path, fun, metadata=metadata) as exporter: + exporter(mx.array([1.0])) + exporter(mx.array([1.0, 2.0])) + _, imported_metadata = mx.import_function(path, return_metadata=True) + self.assertEqual(imported_metadata, metadata) + + # Unsupported value types are rejected + with self.assertRaises(TypeError): + mx.export_function(path, fun, x, metadata={"bad": [1, 2]}) + + # Metadata is not supported with a callback + with self.assertRaises(ValueError): + mx.export_function(lambda x: None, fun, x, metadata=metadata) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/tests/export_import_tests.cpp b/tests/export_import_tests.cpp index ef6a18e199..c8c323ef09 100644 --- a/tests/export_import_tests.cpp +++ b/tests/export_import_tests.cpp @@ -161,3 +161,38 @@ TEST_CASE("test export function on different stream") { // Should make a new stream that we can run computation on eval(import_function(file_path)({array({0, 1, 2})})); } + +TEST_CASE("test export import with metadata") { + std::string file_path = get_temp_file("model.mlxfn"); + + auto fun = [](const std::vector& args) -> std::vector { + return {abs(args[0])}; + }; + + // A value above 2^31 and a non-power-of-two float, to catch narrowing + std::unordered_map metadata = { + {"name", std::string("model")}, + {"params", int64_t(7000000000)}, + {"lr", 0.1}}; + + export_function(file_path, fun, {array({0, 1, 2})}, false, metadata); + + auto imported = import_function(file_path); + CHECK(imported.metadata() == metadata); + CHECK(std::get(imported.metadata().at("name")) == "model"); + CHECK(std::get(imported.metadata().at("params")) == 7000000000); + CHECK(std::get(imported.metadata().at("lr")) == 0.1); + eval(imported({array({0, 1, 2})})); + + // No metadata gives an empty map + export_function(file_path, fun, {array({0, 1, 2})}); + CHECK(import_function(file_path).metadata().empty()); + + // Metadata survives the per-trace header rewrite of a multi-trace export + { + auto fn_exporter = exporter(file_path, fun, false, metadata); + fn_exporter({array({0, 1})}); + fn_exporter({array({0, 1, 2})}); + } + CHECK(import_function(file_path).metadata() == metadata); +}