zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-24 10:51:21 +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() // are currently under a function transformation and the retain_graph()
// function which returns true if we are forced to retain the graph during // function which returns true if we are forced to retain the graph during
// evaluation. // 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::InTracing::grad_counter{0};
int detail::RetainGraph::tracing_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 { struct InTracing {
explicit InTracing(bool dynamic = false, bool grad = false) { explicit InTracing(bool dynamic = false, bool grad = false) {
grad_counter += grad; grad_counter += grad;
trace_stack.push_back({dynamic, grad}); trace_stack().push_back({dynamic, grad});
} }
~InTracing() { ~InTracing() {
grad_counter -= trace_stack.back().second; grad_counter -= trace_stack().back().second;
trace_stack.pop_back(); trace_stack.pop_back();
} }
static bool in_tracing() { static bool in_tracing() {
return !trace_stack.empty(); return !trace_stack().empty();
} }
static bool in_dynamic_tracing() { static bool in_dynamic_tracing() {
// compile is always and only the outer-most transform // 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() { static bool in_grad_tracing() {
@ -43,7 +43,7 @@ struct InTracing {
private: private:
static int grad_counter; static int grad_counter;
static std::vector<std::pair<char, char>> trace_stack; static std::vector<char>& trace_stack();
}; };
struct RetainGraph { struct RetainGraph {