Custom transforms (#1246)

2025-12-16 01:49:05 +08:00 · 2024-07-10 18:00:01 -07:00
parent a3c287354f
commit 5c1fa64fb0
16 changed files with 734 additions and 39 deletions
--- a/mlx/transforms_impl.h
+++ b/mlx/transforms_impl.h
@@ -50,4 +50,20 @@ struct InTracing {
  static int tracing_counter;
 };

+struct RetainGraph {
+  RetainGraph() {
+    tracing_counter++;
+  }
+  ~RetainGraph() {
+    tracing_counter--;
+  }
+
+  static bool retain_graph() {
+    return tracing_counter > 0;
+  }
+
+ private:
+  static int tracing_counter;
+};
+
 } // namespace mlx::core::detail