Compiled should not end in broadcast (#2622)

This commit is contained in:
Angelos Katharopoulos
2025-09-26 13:36:09 -07:00
committed by GitHub
parent d4f4ff3c5e
commit 1a9f820af6

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();
recurse(arr, 0, s, arr.shape());
}