Skip to content

Commit 8310bf8

Browse files
cantoniostensorflower-gardener
authored andcommitted
Fix crash caused by large broadcasting error. There was an int32 overflow.
PiperOrigin-RevId: 475677838
1 parent d328ffb commit 8310bf8

File tree

2 files changed

+9
-2
lines changed

2 files changed

+9
-2
lines changed

Diff for: tensorflow/core/util/bcast.h

+2-2
Original file line numberDiff line numberDiff line change
@@ -134,7 +134,7 @@ BCastList<N>::BCastList(const BCastList::Vec (&x)[N],
134134
typedef BCastList::Vec Vec;
135135

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

@@ -199,7 +199,7 @@ BCastList<N>::BCastList(const BCastList::Vec (&x)[N],
199199
}
200200
Vec output;
201201
bool output_dim_set = false;
202-
int output_dim = -1;
202+
int64_t output_dim = -1;
203203
bool none_is_one = true;
204204
bool set_one = false;
205205
for (int j = 0; j < largest_rank; ++j) {

Diff for: tensorflow/core/util/bcast_test.cc

+7
Original file line numberDiff line numberDiff line change
@@ -375,6 +375,13 @@ TEST(BCastTest, Basic_Tensor_Scalar) {
375375
"[11,7,5,3,2]"
376376
"[11,7,5,3,2]"
377377
"[0,1,2,3,4][]");
378+
379+
// int32 edge-case:
380+
EXPECT_EQ(BCast({1, 2147483648}, {1}),
381+
"[2147483648][1][1][2147483648]"
382+
"[2147483648]"
383+
"[1,2147483648]"
384+
"[0][0,1]");
378385
}
379386

380387
TEST(BCastTest, Basic_Tensor_With_DimSize_1_Scalar) {

0 commit comments

Comments
 (0)