Replace faulty overflow check with a builder for `TensorShape`.
Prevents an integer overflow that was not caught before.
PiperOrigin-RevId: 415381595
Change-Id: I76585ddedc912bd9f4a390aeafa8e2ced1a28863
diff --git a/tensorflow/core/kernels/sparse_tensors_map_ops.cc b/tensorflow/core/kernels/sparse_tensors_map_ops.cc
index 5fa6907..b486a2b 100644
--- a/tensorflow/core/kernels/sparse_tensors_map_ops.cc
+++ b/tensorflow/core/kernels/sparse_tensors_map_ops.cc
@@ -263,22 +263,10 @@
"Rank of input SparseTensor should be > 1, but saw rank: ", rank));
auto input_shape_vec = input_shape->vec<int64_t>();
- int new_num_elements = 1;
- bool overflow_ocurred = false;
- for (int i = 0; i < input_shape_vec.size(); i++) {
- new_num_elements =
- MultiplyWithoutOverflow(new_num_elements, input_shape_vec(i));
- if (new_num_elements < 0) {
- overflow_ocurred = true;
- break;
- }
- }
- OP_REQUIRES(
- context, !overflow_ocurred,
- errors::Internal("Encountered overflow from large input shape."));
-
- TensorShape tensor_input_shape(input_shape_vec);
+ TensorShape tensor_input_shape;
+ OP_REQUIRES_OK(context, TensorShape::BuildTensorShape(input_shape_vec,
+ &tensor_input_shape));
gtl::InlinedVector<int64_t, 8> std_order(rank);
std::iota(std_order.begin(), std_order.end(), 0);
SparseTensor input_st;
