zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-25 03:31:17 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Remove static initializer InTracing::trace_stack

Browse Source
This commit is contained in:
Hanyu Deng 2025-04-10 12:00:52 +08:00 committed by Awni Hannun
parent aaf261d128
commit 83d320c2ea
2 changed files with 9 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
mlx/transforms.cpp
View File

@ -42,7 +42,10 @@ class Synchronizer : public Primitive {
// are currently under a function transformation and the retain_graph()
// function which returns true if we are forced to retain the graph during
// evaluation.
std::vector<std::pair<char, char>> detail::InTracing::trace_stack{};
std::vector<char>& detail::InTracing::trace_stack() {
static std::vector<char> trace_stack_;
return trace_stack_;
}
int detail::InTracing::grad_counter{0};
int detail::RetainGraph::tracing_counter{0};

10
mlx/transforms_impl.h
View File

@ -22,19 +22,19 @@ std::vector<array> vmap_replace(
struct InTracing {
explicit InTracing(bool dynamic = false, bool grad = false) {
grad_counter += grad;
trace_stack.push_back({dynamic, grad});
trace_stack().push_back({dynamic, grad});
}
~InTracing() {
grad_counter -= trace_stack.back().second;
grad_counter -= trace_stack().back().second;
trace_stack.pop_back();
}
static bool in_tracing() {
return !trace_stack.empty();
return !trace_stack().empty();
}
static bool in_dynamic_tracing() {
// compile is always and only the outer-most transform
return in_tracing() && trace_stack.front().first;
return in_tracing() && trace_stack().front().first;
}
static bool in_grad_tracing() {
@ -43,7 +43,7 @@ struct InTracing {
private:
static int grad_counter;
static std::vector<std::pair<char, char>> trace_stack;
static std::vector<char>& trace_stack();
};
struct RetainGraph {