From c10d44181fe173159fd26b8e0be8e66b9e471eec Mon Sep 17 00:00:00 2001
From: Maya Bennett <iarmy0102@gmail.com>
Date: Mon, 20 Jul 2026 21:00:57 -0700
Subject: [PATCH] Fix optimized binary op broadcast rank mismatch

The treat-as-1D path already resizes out to the broadcast target, but the last-dimension and planned broadcast paths still resize it to an operand shape. With leading dimensions of size one, that can change the immutable output rank and fail with InvalidArgument.

Use the broadcast target for those paths and cover both operand orders in add and mul.

Fixes #10959
---
 kernels/optimized/cpu/binary_ops.cpp |  3 +--
 kernels/optimized/cpu/binary_ops.h   |  4 ++--
 kernels/test/op_add_test.cpp         | 23 +++++++++++++++++++++++
 kernels/test/op_mul_test.cpp         | 24 ++++++++++++++++++++++++
 4 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/kernels/optimized/cpu/binary_ops.cpp b/kernels/optimized/cpu/binary_ops.cpp
index 46d8e473051..c3984570402 100644
--- a/kernels/optimized/cpu/binary_ops.cpp
+++ b/kernels/optimized/cpu/binary_ops.cpp
@@ -32,10 +32,9 @@ std::optional<BroadcastElementwisePlan> plan_broadcast_elementwise(
     plan.lhs = &a;
     plan.rhs = &b;
   }
-  auto error = resize_tensor(out, plan.lhs->sizes());
   ET_KERNEL_CHECK_MSG(
       ctx,
-      error == Error::Ok,
+      resize_to_broadcast_target_size(a, b, out) == Error::Ok,
       InvalidArgument,
       std::nullopt,
       "Failed to resize output tensor.");
diff --git a/kernels/optimized/cpu/binary_ops.h b/kernels/optimized/cpu/binary_ops.h
index adbce351f65..063024e9720 100644
--- a/kernels/optimized/cpu/binary_ops.h
+++ b/kernels/optimized/cpu/binary_ops.h
@@ -11,6 +11,7 @@
 #include <executorch/kernels/optimized/vec/functional.h>
 #include <executorch/kernels/portable/cpu/scalar_utils.h>
 #include <executorch/kernels/portable/cpu/util/broadcast_indexes_range.h>
+#include <executorch/kernels/portable/cpu/util/broadcast_util.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
 
 #include <optional>
@@ -189,10 +190,9 @@ Tensor& handle_last_dim_broadcast_elementwise(
     lhs = &a;
     rhs = &b;
   }
-  auto error = resize_tensor(out, lhs->sizes());
   ET_KERNEL_CHECK_MSG(
       ctx,
-      error == Error::Ok,
+      resize_to_broadcast_target_size(a, b, out) == Error::Ok,
       InvalidArgument,
       out,
       "Failed to resize output tensor.");
diff --git a/kernels/test/op_add_test.cpp b/kernels/test/op_add_test.cpp
index bf795d3c47a..947bae5fe50 100644
--- a/kernels/test/op_add_test.cpp
+++ b/kernels/test/op_add_test.cpp
@@ -604,6 +604,29 @@ TEST_F(OpAddOutKernelTest, BroadcastBToA) {
   EXPECT_TENSOR_CLOSE(op_add_out(a, b, 1.0, out), expected);
 }
 
+TEST_F(OpAddOutKernelTest, BroadcastLastDimRankMismatch) {
+  TensorFactory<ScalarType::Float> tf;
+  Tensor a = tf.make({3, 4}, /*data=*/{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+  Tensor b = tf.make({1, 3, 1}, /*data=*/{10, 20, 30});
+  Tensor out = tf.zeros({1, 3, 4});
+
+  Tensor expected = tf.make(
+      {1, 3, 4}, /*data=*/{11, 12, 13, 14, 25, 26, 27, 28, 39, 40, 41, 42});
+  EXPECT_TENSOR_CLOSE(op_add_out(a, b, 1.0, out), expected);
+  EXPECT_TENSOR_CLOSE(op_add_out(b, a, 1.0, out), expected);
+}
+
+TEST_F(OpAddOutKernelTest, Broadcast2dBy1dRankMismatch) {
+  TensorFactory<ScalarType::Float> tf;
+  Tensor a = tf.make({2, 3}, /*data=*/{1, 2, 3, 4, 5, 6});
+  Tensor b = tf.make({1, 1, 3}, /*data=*/{10, 20, 30});
+  Tensor out = tf.zeros({1, 2, 3});
+
+  Tensor expected = tf.make({1, 2, 3}, /*data=*/{11, 22, 33, 14, 25, 36});
+  EXPECT_TENSOR_CLOSE(op_add_out(a, b, 1.0, out), expected);
+  EXPECT_TENSOR_CLOSE(op_add_out(b, a, 1.0, out), expected);
+}
+
 //
 // Death Tests
 //
diff --git a/kernels/test/op_mul_test.cpp b/kernels/test/op_mul_test.cpp
index e4121651b5a..e9b579a34d9 100644
--- a/kernels/test/op_mul_test.cpp
+++ b/kernels/test/op_mul_test.cpp
@@ -923,3 +923,27 @@ TEST_F(OpMulOutTest, BroadcastDimensionMismatchWithDifferentTypes) {
     EXPECT_TENSOR_EQ(result, expected);
   }
 }
+
+TEST_F(OpMulOutTest, BroadcastLastDimRankMismatch) {
+  TensorFactory<ScalarType::Float> tf;
+  Tensor a = tf.make({3, 4}, /*data=*/{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+  Tensor b = tf.make({1, 3, 1}, /*data=*/{10, 20, 30});
+  Tensor out = tf.zeros({1, 3, 4});
+
+  Tensor expected = tf.make(
+      {1, 3, 4},
+      /*data=*/{10, 20, 30, 40, 100, 120, 140, 160, 270, 300, 330, 360});
+  EXPECT_TENSOR_CLOSE(op_mul_out(a, b, out), expected);
+  EXPECT_TENSOR_CLOSE(op_mul_out(b, a, out), expected);
+}
+
+TEST_F(OpMulOutTest, Broadcast2dBy1dRankMismatch) {
+  TensorFactory<ScalarType::Float> tf;
+  Tensor a = tf.make({2, 3}, /*data=*/{1, 2, 3, 4, 5, 6});
+  Tensor b = tf.make({1, 1, 3}, /*data=*/{10, 20, 30});
+  Tensor out = tf.zeros({1, 2, 3});
+
+  Tensor expected = tf.make({1, 2, 3}, /*data=*/{10, 40, 90, 40, 100, 180});
+  EXPECT_TENSOR_CLOSE(op_mul_out(a, b, out), expected);
+  EXPECT_TENSOR_CLOSE(op_mul_out(b, a, out), expected);
+}