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

Fix leak for multi-output primitives which are never detached (#1059)

Browse Source 
* fix multi output leak

* ignore arrays that will be detached

* add some comments

* stray print
This commit is contained in:
Awni Hannun
2024-05-01 07:31:45 -07:00
committed by GitHub
parent 19bef39f5c
commit 7f7b9662ea
5 changed files with 59 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
mlx/transforms.cpp
View File

@@ -246,7 +246,7 @@ std::pair<std::vector<array>, std::vector<array>> vjp(
return;
}
a.set_tracer(false);
for (auto s : a.siblings()) {
for (auto& s : a.siblings()) {
s.set_tracer(false);
cache.insert(s.id());
}
@@ -403,7 +403,7 @@ std::pair<std::vector<array>, std::vector<array>> jvp(
return;
}
a.set_tracer(false);
for (auto s : a.siblings()) {
for (auto& s : a.siblings()) {
s.set_tracer(false);
cache.insert(s.id());
}