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 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 #include #include +#include #include #include @@ -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 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 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 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 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); +}