Include tensor shapes in get_broadcast_target_size error message

kernels/portable/cpu/util/broadcast_util.cpp

@@ -213,10 +213,20 @@ ET_NODISCARD Error get_broadcast_target_size(
     Tensor::SizesType* out_sizes,
     const size_t out_sizes_len,
     size_t* out_dim) {
-  ET_CHECK_OR_RETURN_ERROR(
-      tensors_are_broadcastable_between(a_size, b_size),
-      InvalidArgument,
-      "Two input tensors should be broadcastable.\n");
+  if (!tensors_are_broadcastable_between(a_size, b_size)) {
+    const auto a_shape_str = tensor_shape_to_c_string(
+        executorch::runtime::Span<const Tensor::SizesType>(
+            a_size.data(), a_size.size()));
+    const auto b_shape_str = tensor_shape_to_c_string(
+        executorch::runtime::Span<const Tensor::SizesType>(
+            b_size.data(), b_size.size()));
+    ET_LOG(
+        Error,
+        "Two input tensors should be broadcastable but got shapes %s and %s.",
+        a_shape_str.data(),
+        b_shape_str.data());
+    return executorch::runtime::Error::InvalidArgument;
+  }
 
   auto a_dim = a_size.size();
   auto b_dim = b_size.size();

kernels/portable/cpu/util/test/broadcast_test.cpp

@@ -129,6 +129,15 @@ TEST(BroadcastUtilTest, GetBroadcastTargetSize) {
   EXPECT_TRUE(
       ArrayRef<Tensor::SizesType>(expected_output_size, expected_output_dim)
           .equals(ArrayRef<Tensor::SizesType>({5, 2, 2})));
+
+  Tensor c = tf.zeros({4, 5});
+  err = get_broadcast_target_size(
+      a,
+      c,
+      expected_output_size,
+      torch::executor::kTensorDimensionLimit,
+      &expected_output_dim);
+  EXPECT_EQ(err, torch::executor::Error::InvalidArgument);
 }
 
 size_t linearize_indexes(size_t* indexes, size_t indexes_len, const Tensor& t) {