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

18
mlx/array.h
View File

@@ -261,22 +261,16 @@ class array {
return array_desc_->siblings;
};
/** The array's siblings. */
std::vector<array>& siblings() {
return array_desc_->siblings;
};
void set_siblings(std::vector<array> siblings, uint16_t position) {
array_desc_->siblings = std::move(siblings);
array_desc_->position = position;
}
/** The i-th output of the array's primitive. */
const array& output(int i) const {
if (i == array_desc_->position) {
return *this;
} else if (i < array_desc_->position) {
return siblings()[i];
} else {
return siblings()[i + 1];
}
};
/** The outputs of the array's primitive (i.e. this array and
* its siblings) in the order the primitive expects. */
std::vector<array> outputs() const {
@@ -386,6 +380,8 @@ class array {
array_desc_ = other.array_desc_;
}
~array();
private:
// Initialize the arrays data
template <typename It>