Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix crash caused by large broadcasting error. There was an int32 over…
…flow.

PiperOrigin-RevId: 475677838
  • Loading branch information
cantonios authored and tensorflower-gardener committed Sep 20, 2022
1 parent d328ffb commit 8310bf8
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 2 deletions.
4 changes: 2 additions & 2 deletions tensorflow/core/util/bcast.h
Expand Up @@ -134,7 +134,7 @@ BCastList<N>::BCastList(const BCastList::Vec (&x)[N],
typedef BCastList::Vec Vec;

// Safely multiplies dimensions taking into account symbolic shapes.
auto mul_dims = [](int64_t dim1, int64_t dim2) -> int64 {
auto mul_dims = [](int64_t dim1, int64_t dim2) -> int64_t {
return dim1 != 0 && dim2 != 0 && (dim1 < 0 || dim2 < 0) ? -1 : dim1 * dim2;
};

Expand Down Expand Up @@ -199,7 +199,7 @@ BCastList<N>::BCastList(const BCastList::Vec (&x)[N],
}
Vec output;
bool output_dim_set = false;
int output_dim = -1;
int64_t output_dim = -1;
bool none_is_one = true;
bool set_one = false;
for (int j = 0; j < largest_rank; ++j) {
Expand Down
7 changes: 7 additions & 0 deletions tensorflow/core/util/bcast_test.cc
Expand Up @@ -375,6 +375,13 @@ TEST(BCastTest, Basic_Tensor_Scalar) {
"[11,7,5,3,2]"
"[11,7,5,3,2]"
"[0,1,2,3,4][]");

// int32 edge-case:
EXPECT_EQ(BCast({1, 2147483648}, {1}),
"[2147483648][1][1][2147483648]"
"[2147483648]"
"[1,2147483648]"
"[0][0,1]");
}

TEST(BCastTest, Basic_Tensor_With_DimSize_1_Scalar) {
Expand Down

0 comments on commit 8310bf8

Please sign in to comment.