Skip to content
Permalink
Browse files Browse the repository at this point in the history
Prevent a null pointer dereference in TFLite.
PiperOrigin-RevId: 370800353
Change-Id: Ic9c9712ce5c6e384c954dcd640a5bd9ff05c9a05
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Apr 28, 2021
1 parent 9c1dc92 commit f837892
Showing 1 changed file with 10 additions and 3 deletions.
13 changes: 10 additions & 3 deletions tensorflow/lite/core/subgraph.cc
Expand Up @@ -1060,10 +1060,17 @@ TfLiteStatus Subgraph::Invoke() {
TF_LITE_ENSURE_STATUS(EnsureTensorDataIsReadable(tensor_index));
}
if (tensor->data.raw == nullptr && tensor->bytes > 0) {
if (registration.builtin_code == kTfLiteBuiltinReshape && i == 1) {
if (registration.builtin_code == kTfLiteBuiltinReshape && i == 1 &&
tensor->dims->size != 1) {
// In general, having a tensor here with no buffer will be an error.
// However, for the reshape operator, the second input tensor is only
// used for the shape, not for the data. Thus, null buffer is ok.
// However, for the reshape operator, the second input tensor is
// sometimes only used for the shape, not for the data. Thus, null
// buffer is ok in this situation.
// The situation where null buffer is not ok for reshape operator is
// only when there are 2 inputs given to the node and the one
// corresponding to the shape (i == 1) is a vector that contains all
// dimensions. See `GetOutputShape()` function in
// `tensorflow/lite/kernels/reshape.cc`
continue;
} else {
// In all other cases, we need to return an error as otherwise we will
Expand Down

0 comments on commit f837892

Please sign in to comment.