Skip to content

Commit 37c01fb

Browse files
isharktensorflower-gardener
authored andcommitted
Fix out of bound error in ReverseSequence Op shape function
PiperOrigin-RevId: 411896080 Change-Id: I7e59a38e2f960886edf2b6c54ed5a84e86a9b193
1 parent 3218043 commit 37c01fb

File tree

1 file changed

+10
-0
lines changed

1 file changed

+10
-0
lines changed

Diff for: tensorflow/core/ops/array_ops.cc

+10
Original file line numberDiff line numberDiff line change
@@ -1653,11 +1653,21 @@ REGISTER_OP("ReverseSequence")
16531653
return errors::InvalidArgument(
16541654
"batch_dim must be < input rank: ", batch_dim, " vs. ", input_rank);
16551655
}
1656+
16561657
if (seq_dim >= input_rank) {
16571658
return errors::InvalidArgument(
16581659
"seq_dim must be < input rank: ", seq_dim, " vs. ", input_rank);
16591660
}
16601661

1662+
// To prevent out of bound access when calling c->Dim(input, batch_dim),
1663+
// batch_dim range [-1 * input rank, input rank) is allowed. However,
1664+
// the op implementation has a stricter bound for batch_dim requiring >= 0
1665+
// value. Thus, perform strict check here.
1666+
if (batch_dim < 0) {
1667+
return errors::InvalidArgument("batch_dim must be >=0, got ",
1668+
batch_dim);
1669+
}
1670+
16611671
DimensionHandle batch_dim_dim = c->Dim(input, batch_dim);
16621672
TF_RETURN_IF_ERROR(
16631673
c->Merge(batch_dim_dim, c->Dim(seq_lens_shape, 0), &batch_dim_dim));

0 commit comments

Comments
 (0)