mlx/examples/export
Awni Hannun 4ba0c24a8f
Export / import functions to / from a file (#1642)
* export and import functions

* refactor + works for few primitives

* nit

* allow primitives with state

* nit

* nit

* simplify serialize / deserialize

* fix for constants

* python bindings

* maybe fix serialize failure case

* add example

* more primitives, training kind of works

* same result for python and c++

* some fixes

* fix export

* template it up

* some simplificatoin

* rebase

* allow kwargs and multiple functions

* exporter

* more primitives for exporting

* deal with endianness

* handle invalid stream

* add docstring
2024-12-24 11:19:13 -08:00
..
CMakeLists.txt Export / import functions to / from a file (#1642) 2024-12-24 11:19:13 -08:00
eval_mlp.cpp Export / import functions to / from a file (#1642) 2024-12-24 11:19:13 -08:00
eval_mlp.py Export / import functions to / from a file (#1642) 2024-12-24 11:19:13 -08:00
README.md Export / import functions to / from a file (#1642) 2024-12-24 11:19:13 -08:00
train_mlp.cpp Export / import functions to / from a file (#1642) 2024-12-24 11:19:13 -08:00
train_mlp.py Export / import functions to / from a file (#1642) 2024-12-24 11:19:13 -08:00

Setup

Install MLX:

pip install mlx>=0.22

Build the C++ examples:

cmake -B build -DCMAKE_BUILD_TYPE=Release
cmake --build build

Run

Eval MLP

Run the Python script to export the eval function:

python eval_mlp.py

Then run the C++ program to import and run the function:

./build/eval_mlp

The Python and C++ programs should output the same result.

Train MLP

Run the Python script to export the model initialization and training functions:

python train_mlp.py

Then run the C++ program to import and run the functions:

./build/train_mlp

The Python and C++ programs should output the same results.