diff --git a/kernels/portable/cpu/util/broadcast_util.cpp b/kernels/portable/cpu/util/broadcast_util.cpp index dcc5a0f8e64..aca56ce7e97 100644 --- a/kernels/portable/cpu/util/broadcast_util.cpp +++ b/kernels/portable/cpu/util/broadcast_util.cpp @@ -213,10 +213,22 @@ 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 ET_UNLIKELY (!tensors_are_broadcastable_between(a_size, b_size)) { +#ifdef ET_LOG_ENABLED + const auto a_shape_str = tensor_shape_to_c_string( + executorch::runtime::Span( + a_size.data(), a_size.size())); + const auto b_shape_str = tensor_shape_to_c_string( + executorch::runtime::Span( + b_size.data(), b_size.size())); +#endif + 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(); diff --git a/kernels/portable/cpu/util/test/broadcast_test.cpp b/kernels/portable/cpu/util/test/broadcast_test.cpp index 23640f84689..679296f112c 100644 --- a/kernels/portable/cpu/util/test/broadcast_test.cpp +++ b/kernels/portable/cpu/util/test/broadcast_test.cpp @@ -129,6 +129,15 @@ TEST(BroadcastUtilTest, GetBroadcastTargetSize) { EXPECT_TRUE( ArrayRef(expected_output_size, expected_output_dim) .equals(ArrayRef({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) {