mlx/examples/export/eval_mlp.cpp
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

26 lines
460 B
C++

// Copyright © 2024 Apple Inc.
#include <mlx/mlx.h>
#include <iostream>
using namespace mlx::core;
int main() {
int batch_size = 8;
int input_dim = 32;
// Make the input
random::seed(42);
auto example_x = random::uniform({batch_size, input_dim});
// Import the function
auto forward = import_function("eval_mlp.mlxfn");
// Call the imported function
auto out = forward({example_x})[0];
std::cout << out << std::endl;
return 0;
}