From ce97af5743d042d338cd195a2086bd48d503e98b Mon Sep 17 00:00:00 2001 From: Alan Kelly Date: Thu, 22 Feb 2024 06:06:34 -0800 Subject: [PATCH] Add dimension to subgraph test tensors PiperOrigin-RevId: 609346048 --- tensorflow/lite/kernels/subgraph_test_util.cc | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tensorflow/lite/kernels/subgraph_test_util.cc b/tensorflow/lite/kernels/subgraph_test_util.cc index 2bd0992909a1b8..0521ae31889fa0 100644 --- a/tensorflow/lite/kernels/subgraph_test_util.cc +++ b/tensorflow/lite/kernels/subgraph_test_util.cc @@ -176,6 +176,13 @@ void AddDynamicUpdateSliceNode(Subgraph* subgraph, int input0, int input1, } } // namespace +void Setup1DTensor(Subgraph* subgraph, int tensor_index, TfLiteType type) { + int dim = 1; + ASSERT_EQ(subgraph->SetTensorParametersReadWrite(tensor_index, type, "", 1, + &dim, {}, false), + kTfLiteOk); +} + void SetupTensor(Subgraph* subgraph, int tensor_index, TfLiteType type) { ASSERT_EQ(subgraph->SetTensorParametersReadWrite(tensor_index, type, "", 0, nullptr, {}, false), @@ -275,7 +282,7 @@ void SubgraphBuilder::BuildOutputNotConsumedSubgraph(Subgraph& subgraph) { ASSERT_EQ(subgraph.SetInputs({kInput0, kInput1, kInput2}), kTfLiteOk); ASSERT_EQ(subgraph.SetOutputs({kOutput0, kOutput1, kConstRhs}), kTfLiteOk); for (int i = 0; i < kTensorCount; ++i) { - SetupTensor(&subgraph, i, kTfLiteInt32); + Setup1DTensor(&subgraph, i, kTfLiteInt32); } // kInput0 --> +---+