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

Use unordered map for kwargs in export/import (#2087)

Browse Source 
* use unordered map for kwargs in export/import

* comment
This commit is contained in:
Awni Hannun
2025-04-21 07:17:22 -07:00
committed by GitHub
parent 70ebc3b598
commit dc4eada7f0
4 changed files with 31 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
mlx/export.cpp
View File

@@ -1,5 +1,6 @@
// Copyright © 2024 Apple Inc.
#include "mlx/export.h"
#include <map>
#include "mlx/compile_impl.h"
#include "mlx/fast_primitives.h"
#include "mlx/primitives.h"
@@ -481,7 +482,9 @@ bool FunctionTable::match(
return false;
}
}
for (auto& [_, in] : kwargs) {
auto sorted_kwargs =
std::map<std::string, array>(kwargs.begin(), kwargs.end());
for (auto& [_, in] : sorted_kwargs) {
if (!match_inputs(in, fun.inputs[i++])) {
return false;
}
@@ -557,7 +560,9 @@ void FunctionExporter::export_function(const Args& args, const Kwargs& kwargs) {
// Flatten the inputs to the function for tracing
std::vector<std::string> kwarg_keys;
auto inputs = args;
for (auto& [k, v] : kwargs) {
auto sorted_kwargs =
std::map<std::string, array>(kwargs.begin(), kwargs.end());
for (auto& [k, v] : sorted_kwargs) {
kwarg_keys.push_back(k);
inputs.push_back(v);
}