Usage and Implementation of InTracing and RetainGraph in MLX (transform_impl.h)

[Zijie-Tian]

I’m exploring the internal implementation of graph transformations and autograd mechanics in the MLX framework, specifically focusing on the transform_impl.h file

mlx/mlx/transforms_impl.h

Lines 22 to 74 in 2770a10

	struct InTracing {
	explicit InTracing(bool dynamic = false) {
	trace_stack.push_back(dynamic);
	}
	~InTracing() {
	trace_stack.pop_back();
	}
	static bool in_tracing() {
	return !trace_stack.empty();
	}
	static bool in_dynamic_tracing() {
	// compile is always and only the outer-most transform
	return in_tracing() && trace_stack.front();
	}
	private:
	static std::vector<char> trace_stack;
	};
	struct RetainGraph {
	RetainGraph() {
	tracing_counter++;
	}
	~RetainGraph() {
	tracing_counter--;
	}
	static bool retain_graph() {
	return tracing_counter > 0;
	}
	private:
	static int tracing_counter;
	};
	/** Return true if we are currently performing a function transformation in
	* order to keep the graph when evaluating tracer arrays. */
	inline bool in_tracing() {
	return detail::InTracing::in_tracing();
	}
	/** Return true if we are in a dynamic (shapeless) trace used for compiling or
	* exporting graphs with dynamic shapes. */
	inline bool in_dynamic_tracing() {
	return detail::InTracing::in_dynamic_tracing();
	}
	inline bool retain_graph() {
	return detail::RetainGraph::retain_graph();
	}
	} // namespace mlx::core::detail

I guess these are tracers used to print compute graphs. So,

• How should InTracing be properly used in the computational graph?
• How do they work?
• Any other methods can be used to debug/visualize a compute graph?

[awni]

If you want to visualize a compute graph use export_to_dot. There are other function sin that file which can help print the graph.

The InTracing struct doesn't have to do with printing graphs. That is more so an operation / eval know if they are being called during a transformation.

[Zijie-Tian]

Thank you ! @awni