From 1a9f820af691189053a3570f7ef05a50f54ccccc Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Fri, 26 Sep 2025 13:36:09 -0700 Subject: [PATCH] Compiled should not end in broadcast (#2622) --- mlx/compile.cpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/mlx/compile.cpp b/mlx/compile.cpp index 91743ec04..90eeaa95a 100644 --- a/mlx/compile.cpp +++ b/mlx/compile.cpp @@ -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()); }