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

Compile now can attach arbitrary data to an entry (#2634)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2025-09-30 13:33:27 -07:00
committed by GitHub
parent dc371ae7a5
commit eb24267b56
5 changed files with 130 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
mlx/export.cpp
View File

@@ -579,11 +579,11 @@ void FunctionExporter::export_function(const Args& args, const Kwargs& kwargs) {
for (auto& k : kwarg_keys) {
kwargs.insert({k, *it++});
}
return fun(args, kwargs);
return detail::ArraysAndExtra{fun(args, kwargs), nullptr};
};
// Trace to build the graph
auto [trace_inputs, trace_outputs] =
auto [trace_inputs, trace_outputs, extra] =
detail::compile_trace(flat_fun, inputs, ftable->shapeless);
// DFS the graph and get the tape