[WIP] Export / import functions to / from a file

[awni]

Adds export_function and import_function so that we can save and load functions from a file. Makes it possible to use functions written in one language from another language (e.g. Python -> C++).

Basically works like so:

In Python:

# Note, the model parameters are saved in the export function
# An alternative is to make them inputs to forward
def forward(x):
    return model(x)

example_x = mx.zeros(shape=(batch_size, input_dim))

# Export to file using example input
mx.export_function("model.mlxfn", forward, example_x)

Then in C++, for example:

  auto example_x = random::uniform({batch_size, input_dim});

  // Import the function
  auto forward = import_function("model.mlxfn");

  // Call the imported function
  auto out = forward({example_x})[0];

Some notes on the implementation:

• Reuses a lot of the compile infrastructure which simplifies things dramatically
• The serialization of everything is mostly decoupled from the rest of the code and kept in the export.cpp
• Serializing primitives that have member variables requires some way of accessing them. The API is not opinionated about this (so the primitive interface didn't change at all).. but the convention I'm using is to have a state which returns the data to save
• Likely can use templates / preprocessor to reduce more boiler-plate from some of serialization code in export.cpp. But didn't want to obfuscate / over engineer it much yet until getting some input.

[angeloskath]

This is massively cool. I 'll get to reviewing asap!

[awni]

This change is worth commenting on:

• Previously if you loaded arrays from a file inside a compiled function then every call to the function would reload from the file.
• Now only the first call to the function loads and after that the loaded arrays become constants in the tape

This seems better to me.. though perhaps that is debatable. It is also used by import_function which makes Load primitives for constants and so we get lazy loading even with import_function which is pretty nice.

Lmk thoughts.. I can switch it so compile doesn't force load (more flexible but more dangerous).

[awni]

Another question: should import_function return metadata? I can see how it would be useful to get say the shapes and/or dtypes of the inputs, maybe the MLX version, etc in a dict of metadata. Can also wait and see and provide an overload / a return_metadata flag in the future.

[awni]

Unused 🤷‍♂️ ..