[Feature request] Support saving meta data to mlxfn files and reading them back

[csukuangfj]

Motivation

mlx supports exporting a model with its parameters into an mlxfn file and later loads it into other programming languages, e.g., Python, C++, without the need for the source code of the model, which is very similar to torchscript.

That is a very nice feature.

However, mlx does not support saving meta data into the mlxfn file and later reading it from the code.

Model meta data is a very useful feature. For instance, there are many parakeet models in the mlx-community for ASR. Some of the model uses feature dim 128, while others use 80. Some of the model has 2 LSTM layers in the decoder, while others have only 1 layer.

The feature dim info is required to initialize the feature extractor, while the LSTM info is required to initialize the state array that are later sent to the decoder model.

If we can save such information in the mlxfn file and later read them from the code, it can greatly simplify how users use the model. They are not forced to provide the information about the model files.

onnx supports saving meta data to *.onnx files. See
https://onnxruntime.ai/docs/api/python/auto_examples/plot_metadata.html

There are APIs in onnxruntime to read the meta data.

---

To give you a concrete example, we have converted some parakeet models to mxlfn format
https://github.com/k2-fsa/sherpa-mlx/releases
which can be used in sherpa-mlx.

The problem is that different parakeet models need to be handled differently in the code, e.g., feature dim, decoder LSTM state, etc. But such information is not available from the mlxfn files.

[awni]

We could do something like what we do with safetensors / gguf. The metadata could be a simple dictionary of string to Union[str, int, float] or something like that. And you can load it by requesting return_metadata. That seems pretty reasonable.

That said it also seems pretty easy and way more flexible to serialize your config to a json. Is the downside that you have to manage two files for each model?

[csukuangfj]

Thank you for the quick reply.

We could do something like what we do with

Yes, that would be very nice.

---

to serialize your config to a json

Yes, that is also doable.

Is the downside that you have to manage two files for each model?

Yes.

• As "less code leads to fewer bugs", fewer files can reduce the chance of making mistakes on the user side. E.g., they won't use a wrong json file from a different model.
• You don't need to use a json library to parse the json file or don't need to write code to parse the json file.
• There would be one fewer member in the struct of the public API.

---

By the way, the main usage is in C++, which would later provide bindings for other programming languages, like we do in sherpa-onnx.

[Ritesh102000]

Opened #3682 for this. It adds an optional metadata dict to export_function / exporter that gets written into the .mlxfn file, plus a return_metadata flag on import_function to read it back — same idea as save_safetensors / load. I used string keys/values for consistency with the existing metadata API; let me know if you'd rather support Union[str, int, float] values instead.