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 race in multi-stream eval (#911)

Browse Source 
* maybe fix race

* comment
This commit is contained in:
Awni Hannun
2024-03-26 16:36:36 -07:00
committed by GitHub
parent 29221fa238
commit dc175f08d3
2 changed files with 16 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
mlx/array.h
View File

@@ -256,6 +256,17 @@ class array {
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 {