diff --git a/tensorflow/lite/kernels/kernel_util.h b/tensorflow/lite/kernels/kernel_util.h index 6bd6bb1c7ed1141ab50e5ad8e16248366b1ce66f..8fae5740e50209166d42d38f142fba34475480fd 100644 --- a/tensorflow/lite/kernels/kernel_util.h +++ b/tensorflow/lite/kernels/kernel_util.h @@ -30,27 +30,49 @@ inline int SizeOfDimension(const TfLiteTensor* t, int dim) { } inline const TfLiteTensor* GetInput(const TfLiteContext* context, const TfLiteNode* node, int index) { - return &context->tensors[node->inputs->data[index]]; + const int tensor_index = node->inputs->data[index]; + if (tensor_index < 0) { + return nullptr; + } + return &context->tensors[tensor_index]; } // Note: You must check if result is not null: // TfLiteTensor* my_tensor = GetVariableInput(context, node, kMyTensorIdx); // TF_LITE_ENSURE(context, my_tensor != nullptr); inline TfLiteTensor* GetVariableInput(TfLiteContext* context, const TfLiteNode* node, int index) { - TfLiteTensor* tensor = &context->tensors[node->inputs->data[index]]; + const int tensor_index = node->inputs->data[index]; + if (tensor_index < 0) { + return nullptr; + } + TfLiteTensor* tensor = &context->tensors[tensor_index]; +>>>>>>> d8f8236c29 ([tflite] Test for `kTfLiteOptionalTensor` in `GetInput`.) return (tensor->is_variable) ? tensor : nullptr; } inline TfLiteTensor* GetOutput(TfLiteContext* context, const TfLiteNode* node, int index) { - return &context->tensors[node->outputs->data[index]]; + const int tensor_index = node->outputs->data[index]; + if (tensor_index < 0) { + return nullptr; + } + return &context->tensors[tensor_index]; } inline TfLiteTensor* GetTemporary(TfLiteContext* context, const TfLiteNode* node, int index) { - return &context->tensors[node->temporaries->data[index]]; + const int tensor_index = node->temporaries->data[index]; + if (tensor_index < 0) { + return nullptr; + } + return &context->tensors[tensor_index]; } + inline const TfLiteTensor* GetIntermediates(TfLiteContext* context, const TfLiteNode* node, int index) { - return &context->tensors[node->intermediates->data[index]]; + const int tensor_index = node->intermediates->data[index]; + if (tensor_index < 0) { + return nullptr; + } + return &context->tensors[tensor_index]; } inline int NumInputs(const TfLiteNode* node) { return node->inputs->size; } inline int NumOutputs(const TfLiteNode* node) { return node->outputs->size; }