zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-19 16:48:10 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Compiled should not end in broadcast (#2622)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2025-09-26 13:36:09 -07:00
committed by GitHub
parent d4f4ff3c5e
commit 1a9f820af6
1 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
mlx/compile.cpp
View File

@@ -727,7 +727,11 @@ void compile_fuse(
} }
}; };
if (arr.has_primitive()) { // This will be the result of the fused operation so it needs
// a) to not be already computed ie have a primitive
// b) that primitive to not be a broadcast since it will unnecessarily
// cast to a contiguous array potentially blowing up memory
if (arr.has_primitive() && !is_broadcast(arr.primitive())) {
Stream s = arr.primitive().stream(); Stream s = arr.primitive().stream();
recurse(arr, 0, s, arr.shape()); recurse(arr, 0, s, arr.shape());
} }