Merge pull request #49701 from geetachavan1/cherrypicks_ZNR2C

Prevent overflow in sparse op
tensorflow · May 26, 2021 · e4de6e9 · e4de6e9
2 parents 01e36a0 + 5788883
commit e4de6e9
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/tensorflow/core/kernels/sparse_split_op.cc b/tensorflow/core/kernels/sparse_split_op.cc
@@ -63,11 +63,18 @@ class SparseSplitOp : public OpKernel {
                                 input_shape.vec<int64>()(split_dim), "), got ",
                                 num_split_));
 
+    // Prevent overflow by constructing the dense shape separately
+    TensorShape dense_shape;
+    const auto input_shape_flat = input_shape.flat<int64>();
+    for (int i = 0; i < input_shape.NumElements(); i++) {
+      OP_REQUIRES_OK(context,
+                     dense_shape.AddDimWithStatus(input_shape_flat(i)));
+    }
+
     sparse::SparseTensor sparse_tensor;
     OP_REQUIRES_OK(context,
-                   sparse::SparseTensor::Create(
-                       input_indices, input_values,
-                       TensorShape(input_shape.vec<int64>()), &sparse_tensor));
+                   sparse::SparseTensor::Create(input_indices, input_values,
+                                                dense_shape, &sparse_tensor));
 
     std::vector<sparse::SparseTensor> outputs;
     OP_REQUIRES_OK(context,