Skip to content

Commit b51b82f

Browse files
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
1 parent b68e5c4 commit b51b82f

File tree

1 file changed

+15
-2
lines changed

1 file changed

+15
-2
lines changed

tensorflow/core/kernels/sparse_tensors_map_ops.cc

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -231,16 +231,29 @@ class AddManySparseToTensorsMapOp : public SparseTensorAccessingOp {
231231
errors::InvalidArgument(
232232
"Input indices should be a matrix but received shape ",
233233
input_indices->shape().DebugString()));
234-
235234
OP_REQUIRES(context, TensorShapeUtils::IsVector(input_values->shape()),
236235
errors::InvalidArgument(
237236
"Input values should be a vector but received shape ",
238237
input_values->shape().DebugString()));
239-
240238
OP_REQUIRES(context, TensorShapeUtils::IsVector(input_shape->shape()),
241239
errors::InvalidArgument(
242240
"Input shape should be a vector but received shape ",
243241
input_shape->shape().DebugString()));
242+
OP_REQUIRES(
243+
context,
244+
input_values->shape().dim_size(0) == input_indices->shape().dim_size(0),
245+
errors::InvalidArgument(
246+
"Number of values must match first dimension of indices. ", "Got ",
247+
input_values->shape().dim_size(0),
248+
" values, indices shape: ", input_indices->shape().DebugString()));
249+
OP_REQUIRES(
250+
context,
251+
input_shape->shape().dim_size(0) == input_indices->shape().dim_size(1),
252+
errors::InvalidArgument(
253+
"Number of dimensions must match second dimension of indices. ",
254+
"Got ", input_shape->shape().dim_size(0),
255+
" dimensions, indices shape: ",
256+
input_indices->shape().DebugString()));
244257

245258
int rank = input_shape->NumElements();
246259

0 commit comments

Comments
 (0)