zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-20 01:18:12 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

MLX in C++ example (#1736)

Browse Source 
* MLX in C++ example

* nits

* fix docs
This commit is contained in:
Awni Hannun
2025-01-02 19:09:04 -08:00
committed by GitHub
parent 8544b42007
commit c9d30aa6ac
13 changed files with 242 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
examples/export/train_mlp.cpp
View File

@@ -3,22 +3,22 @@
#include <mlx/mlx.h>
#include <iostream>
using namespace mlx::core;
namespace mx = mlx::core;
int main() {
int batch_size = 8;
int input_dim = 32;
int output_dim = 10;
auto state = import_function("init_mlp.mlxfn")({});
auto state = mx::import_function("init_mlp.mlxfn")({});
// Make the input
random::seed(42);
auto example_X = random::normal({batch_size, input_dim});
auto example_y = random::randint(0, output_dim, {batch_size});
mx::random::seed(42);
auto example_X = mx::random::normal({batch_size, input_dim});
auto example_y = mx::random::randint(0, output_dim, {batch_size});
// Import the function
auto step = import_function("train_mlp.mlxfn");
auto step = mx::import_function("train_mlp.mlxfn");
// Call the imported function
for (int it = 0; it < 100; ++it) {