Skip to content

Commit

Permalink
Relocate output tensor dims for GATHER_ND
Browse files Browse the repository at this point in the history
@tensorflow/micro

Move the output tensor dims data for GATHER_ND using the CreateWritableTensorDimsWithCopy utility method.

bug=fixes `gather_nd`: prepare stage tries to access model `.rodata` tensorflow#1723
  • Loading branch information
ddavis-2015 committed Feb 9, 2023
1 parent 6343983 commit 71c4a36
Showing 1 changed file with 7 additions and 0 deletions.
7 changes: 7 additions & 0 deletions tensorflow/lite/micro/kernels/gather_nd.cc
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,13 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
// Assign to output the input type.
output->type = params->type;

// The tensor output dims must be relocated
// from the FlatBuffer to the persistant storage arena.
TfLiteEvalTensor* output_eval =
tflite::micro::GetEvalOutput(context, node, kOutputTensor);
TF_LITE_ENSURE_OK(context, tflite::micro::CreateWritableTensorDimsWithCopy(
context, output, output_eval));

// TFLM gather_nd does not create the output tensor, but it needs to ensure
// that the output shape is correct. The result shape is
// indices.shape[:-1] + params.shape[indices.shape[-1]:]
Expand Down

0 comments on commit 71c4a36

Please sign in to comment.