Skip to content

Commit a68f680

Browse files
Replace faulty overflow check with a builder for TensorShape.
Prevents an integer overflow that was not caught before. PiperOrigin-RevId: 415381595 Change-Id: I76585ddedc912bd9f4a390aeafa8e2ced1a28863
1 parent e8f4be7 commit a68f680

File tree

1 file changed

+3
-15
lines changed

1 file changed

+3
-15
lines changed

Diff for: tensorflow/core/kernels/sparse_tensors_map_ops.cc

+3-15
Original file line numberDiff line numberDiff line change
@@ -263,22 +263,10 @@ class AddManySparseToTensorsMapOp : public SparseTensorAccessingOp {
263263
"Rank of input SparseTensor should be > 1, but saw rank: ", rank));
264264

265265
auto input_shape_vec = input_shape->vec<int64_t>();
266-
int new_num_elements = 1;
267-
bool overflow_ocurred = false;
268-
for (int i = 0; i < input_shape_vec.size(); i++) {
269-
new_num_elements =
270-
MultiplyWithoutOverflow(new_num_elements, input_shape_vec(i));
271-
if (new_num_elements < 0) {
272-
overflow_ocurred = true;
273-
break;
274-
}
275-
}
276-
277-
OP_REQUIRES(
278-
context, !overflow_ocurred,
279-
errors::Internal("Encountered overflow from large input shape."));
280266

281-
TensorShape tensor_input_shape(input_shape_vec);
267+
TensorShape tensor_input_shape;
268+
OP_REQUIRES_OK(context, TensorShape::BuildTensorShape(input_shape_vec,
269+
&tensor_input_shape));
282270
gtl::InlinedVector<int64_t, 8> std_order(rank);
283271
std::iota(std_order.begin(), std_order.end(), 0);
284272
SparseTensor input_st;

0 commit comments

Comments
 (0)