From a8ea75afb810535705341175caa87a3fc2f48bd0 Mon Sep 17 00:00:00 2001
From: Angelos Katharopoulos <a_katharopoulos@apple.com>
Date: Tue, 31 Mar 2026 00:36:18 -0700
Subject: [PATCH 1/3] Fix use after move

---
 mlx/ops.cpp | 101 +++++++++++++++++++++++++++++++---------------------
 1 file changed, 60 insertions(+), 41 deletions(-)

diff --git a/mlx/ops.cpp b/mlx/ops.cpp
index ef792cd6f4..d60e2eab78 100644
--- a/mlx/ops.cpp
+++ b/mlx/ops.cpp
@@ -1658,11 +1658,9 @@ std::vector<array> broadcast_arrays(
       if (in.shape() == out_shape) {
         outputs.push_back(in);
       } else {
-        outputs.push_back(array(
-            out_shape,
-            in.dtype(),
-            std::make_shared<Broadcast>(to_stream(s), out_shape),
-            {in}));
+        auto prim = std::make_shared<Broadcast>(to_stream(s), out_shape);
+        outputs.push_back(
+            array(std::move(out_shape), in.dtype(), std::move(prim), {in}));
       }
     }
     return outputs;
@@ -1766,17 +1764,20 @@ std::pair<array, array> broadcast_arrays(
 array equal(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   auto dtype = promote_types(a.dtype(), b.dtype());
   auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
-  auto& shape = inputs[0].shape();
+  auto shape = inputs[0].shape();
   return array(
-      shape, bool_, std::make_shared<Equal>(to_stream(s)), std::move(inputs));
+      std::move(shape),
+      bool_,
+      std::make_shared<Equal>(to_stream(s)),
+      std::move(inputs));
 }
 
 array not_equal(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   auto dtype = promote_types(a.dtype(), b.dtype());
   auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
-  auto& shape = inputs[0].shape();
+  auto shape = inputs[0].shape();
   return array(
-      shape,
+      std::move(shape),
       bool_,
       std::make_shared<NotEqual>(to_stream(s)),
       std::move(inputs));
@@ -1785,9 +1786,12 @@ array not_equal(const array& a, const array& b, StreamOrDevice s /* = {} */) {
 array greater(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   auto dtype = promote_types(a.dtype(), b.dtype());
   auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
-  auto& shape = inputs[0].shape();
+  auto shape = inputs[0].shape();
   return array(
-      shape, bool_, std::make_shared<Greater>(to_stream(s)), std::move(inputs));
+      std::move(shape),
+      bool_,
+      std::make_shared<Greater>(to_stream(s)),
+      std::move(inputs));
 }
 
 array greater_equal(
@@ -1796,9 +1800,9 @@ array greater_equal(
     StreamOrDevice s /* = {} */) {
   auto dtype = promote_types(a.dtype(), b.dtype());
   auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
-  auto& shape = inputs[0].shape();
+  auto shape = inputs[0].shape();
   return array(
-      shape,
+      std::move(shape),
       bool_,
       std::make_shared<GreaterEqual>(to_stream(s)),
       std::move(inputs));
@@ -1807,17 +1811,20 @@ array greater_equal(
 array less(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   auto dtype = promote_types(a.dtype(), b.dtype());
   auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
-  auto& shape = inputs[0].shape();
+  auto shape = inputs[0].shape();
   return array(
-      shape, bool_, std::make_shared<Less>(to_stream(s)), std::move(inputs));
+      std::move(shape),
+      bool_,
+      std::make_shared<Less>(to_stream(s)),
+      std::move(inputs));
 }
 
 array less_equal(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   auto dtype = promote_types(a.dtype(), b.dtype());
   auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
-  auto& shape = inputs[0].shape();
+  auto shape = inputs[0].shape();
   return array(
-      shape,
+      std::move(shape),
       bool_,
       std::make_shared<LessEqual>(to_stream(s)),
       std::move(inputs));
@@ -2811,9 +2818,9 @@ array logical_not(const array& a, StreamOrDevice s /* = {} */) {
 array logical_and(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   // Broadcast arrays to a common shape
   auto inputs = broadcast_arrays({astype(a, bool_, s), astype(b, bool_, s)}, s);
-  auto& shape = inputs[0].shape();
+  auto shape = inputs[0].shape();
   return array(
-      shape,
+      std::move(shape),
       bool_,
       std::make_shared<LogicalAnd>(to_stream(s)),
       std::move(inputs));
@@ -2825,9 +2832,9 @@ array operator&&(const array& a, const array& b) {
 array logical_or(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   // Broadcast arrays to a common shape
   auto inputs = broadcast_arrays({astype(a, bool_, s), astype(b, bool_, s)}, s);
-  auto& shape = inputs[0].shape();
+  auto shape = inputs[0].shape();
   return array(
-      shape,
+      std::move(shape),
       bool_,
       std::make_shared<LogicalOr>(to_stream(s)),
       std::move(inputs));
@@ -2845,9 +2852,12 @@ array add(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   auto out_type = promote_types(a.dtype(), b.dtype());
   auto inputs =
       broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s);
-  auto& shape = inputs[0].shape();
+  auto shape = inputs[0].shape();
   return array(
-      shape, out_type, std::make_shared<Add>(to_stream(s)), std::move(inputs));
+      std::move(shape),
+      out_type,
+      std::make_shared<Add>(to_stream(s)),
+      std::move(inputs));
 }
 
 array operator+(const array& a, const array& b) {
@@ -2858,9 +2868,9 @@ array subtract(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   auto out_type = promote_types(a.dtype(), b.dtype());
   auto inputs =
       broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s);
-  auto& shape = inputs[0].shape();
+  auto shape = inputs[0].shape();
   return array(
-      shape,
+      std::move(shape),
       out_type,
       std::make_shared<Subtract>(to_stream(s)),
       std::move(inputs));
@@ -2874,9 +2884,9 @@ array multiply(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   auto out_type = promote_types(a.dtype(), b.dtype());
   auto inputs =
       broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s);
-  auto& shape = inputs[0].shape();
+  auto shape = inputs[0].shape();
   return array(
-      shape,
+      std::move(shape),
       out_type,
       std::make_shared<Multiply>(to_stream(s)),
       std::move(inputs));
@@ -2890,9 +2900,12 @@ array divide(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   auto dtype = at_least_float(promote_types(a.dtype(), b.dtype()));
   auto inputs = broadcast_arrays(
       {astype(a, dtype, s), astype(b, dtype, to_stream(s))}, s);
-  auto& shape = inputs[0].shape();
+  auto shape = inputs[0].shape();
   return array(
-      shape, dtype, std::make_shared<Divide>(to_stream(s)), std::move(inputs));
+      std::move(shape),
+      dtype,
+      std::make_shared<Divide>(to_stream(s)),
+      std::move(inputs));
 }
 array operator/(const array& a, const array& b) {
   return divide(a, b);
@@ -2914,18 +2927,21 @@ array floor_divide(
   }
 
   auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
-  auto& shape = inputs[0].shape();
+  auto shape = inputs[0].shape();
   return array(
-      shape, dtype, std::make_shared<Divide>(to_stream(s)), std::move(inputs));
+      std::move(shape),
+      dtype,
+      std::make_shared<Divide>(to_stream(s)),
+      std::move(inputs));
 }
 
 array remainder(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   auto dtype = promote_types(a.dtype(), b.dtype());
   auto inputs = broadcast_arrays(
       {astype(a, dtype, s), astype(b, dtype, to_stream(s))}, s);
-  auto& shape = inputs[0].shape();
+  auto shape = inputs[0].shape();
   return array(
-      shape,
+      std::move(shape),
       dtype,
       std::make_shared<Remainder>(to_stream(s)),
       std::move(inputs));
@@ -2953,9 +2969,9 @@ array maximum(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   auto out_type = promote_types(a.dtype(), b.dtype());
   auto inputs =
       broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s);
-  auto& shape = inputs[0].shape();
+  auto shape = inputs[0].shape();
   return array(
-      shape,
+      std::move(shape),
       out_type,
       std::make_shared<Maximum>(to_stream(s)),
       std::move(inputs));
@@ -2965,9 +2981,9 @@ array minimum(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   auto out_type = promote_types(a.dtype(), b.dtype());
   auto inputs =
       broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s);
-  auto& shape = inputs[0].shape();
+  auto shape = inputs[0].shape();
   return array(
-      shape,
+      std::move(shape),
       out_type,
       std::make_shared<Minimum>(to_stream(s)),
       std::move(inputs));
@@ -3048,9 +3064,12 @@ array arctan(const array& a, StreamOrDevice s /* = {} */) {
 array arctan2(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   auto dtype = at_least_float(promote_types(a.dtype(), b.dtype()));
   auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
-  auto& shape = inputs[0].shape();
+  auto shape = inputs[0].shape();
   return array(
-      shape, dtype, std::make_shared<ArcTan2>(to_stream(s)), std::move(inputs));
+      std::move(shape),
+      dtype,
+      std::make_shared<ArcTan2>(to_stream(s)),
+      std::move(inputs));
 }
 
 array sinh(const array& a, StreamOrDevice s /* = {} */) {
@@ -3144,9 +3163,9 @@ array logaddexp(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   auto out_type = at_least_float(promote_types(a.dtype(), b.dtype()));
   auto inputs =
       broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s);
-  auto& shape = inputs[0].shape();
+  auto shape = inputs[0].shape();
   return array(
-      shape,
+      std::move(shape),
       out_type,
       std::make_shared<LogAddExp>(to_stream(s)),
       std::move(inputs));

From ab0e1eff2795f029807cc95e3305969a9a7f72d4 Mon Sep 17 00:00:00 2001
From: Angelos Katharopoulos <a_katharopoulos@apple.com>
Date: Tue, 31 Mar 2026 00:58:42 -0700
Subject: [PATCH 2/3] A last one

---
 mlx/ops.cpp | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/mlx/ops.cpp b/mlx/ops.cpp
index d60e2eab78..042a9b8e53 100644
--- a/mlx/ops.cpp
+++ b/mlx/ops.cpp
@@ -4206,10 +4206,12 @@ array conv_transpose_general(
         output_padding[i]; // Adjust with output_padding
   }
 
+  auto ndim = stride.size();
+
   return conv_general(
       /* const array& input = */ input,
       /* const array& weight = */ weight,
-      /* std::vector<int> stride = */ std::vector(stride.size(), 1),
+      /* std::vector<int> stride = */ std::vector(ndim, 1),
       /* std::vector<int> padding_lo = */ std::move(padding_lo),
       /* std::vector<int> padding_hi = */ std::move(padding_hi),
       /* std::vector<int> kernel_dilation = */ std::move(dilation),

From dd7242808f0d5f2100cfe960bdb50ef205694e89 Mon Sep 17 00:00:00 2001
From: Angelos Katharopoulos <a_katharopoulos@apple.com>
Date: Tue, 31 Mar 2026 03:37:50 -0700
Subject: [PATCH 3/3] Remove shape movements

---
 mlx/ops.cpp | 65 +++++++++++++++++++----------------------------------
 1 file changed, 23 insertions(+), 42 deletions(-)

diff --git a/mlx/ops.cpp b/mlx/ops.cpp
index 042a9b8e53..8bc13952cc 100644
--- a/mlx/ops.cpp
+++ b/mlx/ops.cpp
@@ -1658,9 +1658,11 @@ std::vector<array> broadcast_arrays(
       if (in.shape() == out_shape) {
         outputs.push_back(in);
       } else {
-        auto prim = std::make_shared<Broadcast>(to_stream(s), out_shape);
-        outputs.push_back(
-            array(std::move(out_shape), in.dtype(), std::move(prim), {in}));
+        outputs.push_back(array(
+            out_shape,
+            in.dtype(),
+            std::make_shared<Broadcast>(to_stream(s), out_shape),
+            {in}));
       }
     }
     return outputs;
@@ -1766,10 +1768,7 @@ array equal(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
   auto shape = inputs[0].shape();
   return array(
-      std::move(shape),
-      bool_,
-      std::make_shared<Equal>(to_stream(s)),
-      std::move(inputs));
+      shape, bool_, std::make_shared<Equal>(to_stream(s)), std::move(inputs));
 }
 
 array not_equal(const array& a, const array& b, StreamOrDevice s /* = {} */) {
@@ -1777,7 +1776,7 @@ array not_equal(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
   auto shape = inputs[0].shape();
   return array(
-      std::move(shape),
+      shape,
       bool_,
       std::make_shared<NotEqual>(to_stream(s)),
       std::move(inputs));
@@ -1788,10 +1787,7 @@ array greater(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
   auto shape = inputs[0].shape();
   return array(
-      std::move(shape),
-      bool_,
-      std::make_shared<Greater>(to_stream(s)),
-      std::move(inputs));
+      shape, bool_, std::make_shared<Greater>(to_stream(s)), std::move(inputs));
 }
 
 array greater_equal(
@@ -1802,7 +1798,7 @@ array greater_equal(
   auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
   auto shape = inputs[0].shape();
   return array(
-      std::move(shape),
+      shape,
       bool_,
       std::make_shared<GreaterEqual>(to_stream(s)),
       std::move(inputs));
@@ -1813,10 +1809,7 @@ array less(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
   auto shape = inputs[0].shape();
   return array(
-      std::move(shape),
-      bool_,
-      std::make_shared<Less>(to_stream(s)),
-      std::move(inputs));
+      shape, bool_, std::make_shared<Less>(to_stream(s)), std::move(inputs));
 }
 
 array less_equal(const array& a, const array& b, StreamOrDevice s /* = {} */) {
@@ -1824,7 +1817,7 @@ array less_equal(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
   auto shape = inputs[0].shape();
   return array(
-      std::move(shape),
+      shape,
       bool_,
       std::make_shared<LessEqual>(to_stream(s)),
       std::move(inputs));
@@ -2820,7 +2813,7 @@ array logical_and(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   auto inputs = broadcast_arrays({astype(a, bool_, s), astype(b, bool_, s)}, s);
   auto shape = inputs[0].shape();
   return array(
-      std::move(shape),
+      shape,
       bool_,
       std::make_shared<LogicalAnd>(to_stream(s)),
       std::move(inputs));
@@ -2834,7 +2827,7 @@ array logical_or(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   auto inputs = broadcast_arrays({astype(a, bool_, s), astype(b, bool_, s)}, s);
   auto shape = inputs[0].shape();
   return array(
-      std::move(shape),
+      shape,
       bool_,
       std::make_shared<LogicalOr>(to_stream(s)),
       std::move(inputs));
@@ -2854,10 +2847,7 @@ array add(const array& a, const array& b, StreamOrDevice s /* = {} */) {
       broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s);
   auto shape = inputs[0].shape();
   return array(
-      std::move(shape),
-      out_type,
-      std::make_shared<Add>(to_stream(s)),
-      std::move(inputs));
+      shape, out_type, std::make_shared<Add>(to_stream(s)), std::move(inputs));
 }
 
 array operator+(const array& a, const array& b) {
@@ -2870,7 +2860,7 @@ array subtract(const array& a, const array& b, StreamOrDevice s /* = {} */) {
       broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s);
   auto shape = inputs[0].shape();
   return array(
-      std::move(shape),
+      shape,
       out_type,
       std::make_shared<Subtract>(to_stream(s)),
       std::move(inputs));
@@ -2886,7 +2876,7 @@ array multiply(const array& a, const array& b, StreamOrDevice s /* = {} */) {
       broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s);
   auto shape = inputs[0].shape();
   return array(
-      std::move(shape),
+      shape,
       out_type,
       std::make_shared<Multiply>(to_stream(s)),
       std::move(inputs));
@@ -2902,10 +2892,7 @@ array divide(const array& a, const array& b, StreamOrDevice s /* = {} */) {
       {astype(a, dtype, s), astype(b, dtype, to_stream(s))}, s);
   auto shape = inputs[0].shape();
   return array(
-      std::move(shape),
-      dtype,
-      std::make_shared<Divide>(to_stream(s)),
-      std::move(inputs));
+      shape, dtype, std::make_shared<Divide>(to_stream(s)), std::move(inputs));
 }
 array operator/(const array& a, const array& b) {
   return divide(a, b);
@@ -2929,10 +2916,7 @@ array floor_divide(
   auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
   auto shape = inputs[0].shape();
   return array(
-      std::move(shape),
-      dtype,
-      std::make_shared<Divide>(to_stream(s)),
-      std::move(inputs));
+      shape, dtype, std::make_shared<Divide>(to_stream(s)), std::move(inputs));
 }
 
 array remainder(const array& a, const array& b, StreamOrDevice s /* = {} */) {
@@ -2941,7 +2925,7 @@ array remainder(const array& a, const array& b, StreamOrDevice s /* = {} */) {
       {astype(a, dtype, s), astype(b, dtype, to_stream(s))}, s);
   auto shape = inputs[0].shape();
   return array(
-      std::move(shape),
+      shape,
       dtype,
       std::make_shared<Remainder>(to_stream(s)),
       std::move(inputs));
@@ -2971,7 +2955,7 @@ array maximum(const array& a, const array& b, StreamOrDevice s /* = {} */) {
       broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s);
   auto shape = inputs[0].shape();
   return array(
-      std::move(shape),
+      shape,
       out_type,
       std::make_shared<Maximum>(to_stream(s)),
       std::move(inputs));
@@ -2983,7 +2967,7 @@ array minimum(const array& a, const array& b, StreamOrDevice s /* = {} */) {
       broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s);
   auto shape = inputs[0].shape();
   return array(
-      std::move(shape),
+      shape,
       out_type,
       std::make_shared<Minimum>(to_stream(s)),
       std::move(inputs));
@@ -3066,10 +3050,7 @@ array arctan2(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
   auto shape = inputs[0].shape();
   return array(
-      std::move(shape),
-      dtype,
-      std::make_shared<ArcTan2>(to_stream(s)),
-      std::move(inputs));
+      shape, dtype, std::make_shared<ArcTan2>(to_stream(s)), std::move(inputs));
 }
 
 array sinh(const array& a, StreamOrDevice s /* = {} */) {
@@ -3165,7 +3146,7 @@ array logaddexp(const array& a, const array& b, StreamOrDevice s /* = {} */) {
       broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s);
   auto shape = inputs[0].shape();
   return array(
-      std::move(shape),
+      shape,
       out_type,
       std::make_shared<LogAddExp>(to_stream(s)),
       std::move(inputs));