commit b51b82fe65ebace4475e3c54eb089c18a4403f1c
author Mihai Maruseac <mihaimaruseac@google.com> Thu Dec 09 14:32:48 2021 -0800
committer TensorFlower Gardener <gardener@tensorflow.org> Thu Dec 09 14:41:08 2021 -0800
tree 4a45f0b5b7bb1ea761a91e5962ec18cfe4768350
parent b68e5c40bddeacda9b09a7cb694d55e05e94e6be

Add missing validation to `AddManySparseToTensorsMap`.

Sparse tensors have a set of requirements for the 3 components and not all of them were checked.

PiperOrigin-RevId: 415358027
Change-Id: I96cbb672999cd1da772c22fabbd15507e32e12dc

tensorflow/core/kernels/sparse_tensors_map_ops.cc

diff --git a/tensorflow/core/kernels/sparse_tensors_map_ops.cc b/tensorflow/core/kernels/sparse_tensors_map_ops.cc
index 04efed5..5fa6907 100644
--- a/tensorflow/core/kernels/sparse_tensors_map_ops.cc
+++ b/tensorflow/core/kernels/sparse_tensors_map_ops.cc
@@ -231,16 +231,29 @@
                 errors::InvalidArgument(
                     "Input indices should be a matrix but received shape ",
                     input_indices->shape().DebugString()));
-
     OP_REQUIRES(context, TensorShapeUtils::IsVector(input_values->shape()),
                 errors::InvalidArgument(
                     "Input values should be a vector but received shape ",
                     input_values->shape().DebugString()));
-
     OP_REQUIRES(context, TensorShapeUtils::IsVector(input_shape->shape()),
                 errors::InvalidArgument(
                     "Input shape should be a vector but received shape ",
                     input_shape->shape().DebugString()));
+    OP_REQUIRES(
+        context,
+        input_values->shape().dim_size(0) == input_indices->shape().dim_size(0),
+        errors::InvalidArgument(
+            "Number of values must match first dimension of indices. ", "Got ",
+            input_values->shape().dim_size(0),
+            " values, indices shape: ", input_indices->shape().DebugString()));
+    OP_REQUIRES(
+        context,
+        input_shape->shape().dim_size(0) == input_indices->shape().dim_size(1),
+        errors::InvalidArgument(
+            "Number of dimensions must match second dimension of indices. ",
+            "Got ", input_shape->shape().dim_size(0),
+            " dimensions, indices shape: ",
+            input_indices->shape().DebugString()));
 
     int rank = input_shape->NumElements();