Fix bad optimized kernel for add.

Differential Revision: D80914321

Pull Request resolved: #13633

Author: hsharma35

kernels/optimized/cpu/op_add_sub_impl.h

@@ -116,10 +116,9 @@ Tensor& opt_add_sub_out_impl(
 
   if (selected_optimized_path == ElementwiseOptimizedPath::kTreatAs1d) {
     // Resize for dynamic shape
-    auto error = resize_tensor(out, a.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.");

kernels/test/op_add_test.cpp

@@ -591,6 +591,18 @@ TEST_F(OpAddOutKernelTest, BroadcastNDTest) {
   test_broadcast_last_dim<ScalarType::BFloat16>();
 }
 
+TEST_F(OpAddOutKernelTest, BroadcastBToA) {
+  TensorFactory<ScalarType::Float> tf_a;
+  Tensor a = tf_a.make({1, 3}, /*data=*/{1, 2, 3});
+  Tensor b = tf_a.make({1, 1, 3}, /*data=*/{3.2, 1.3, 5.5});
+  // Destination for output of add.
+  Tensor out = tf_a.zeros({1, 1, 3});
+
+  // Check that it matches the expected output.
+  Tensor expected = tf_a.make({1, 1, 3}, /*data=*/{4.2, 3.3, 8.5});
+  EXPECT_TENSOR_CLOSE(op_add_out(a, b, 1.0, out), expected);
+}
+
 //
 // Death Tests
 //