| 307 | } |
| 308 | |
| 309 | void SubgraphBuilder::BuildWhileSubgraph(Subgraph* subgraph) { |
| 310 | const int kInput1 = 0; |
| 311 | const int kInput2 = 1; |
| 312 | const int kOutput1 = 2; |
| 313 | const int kOutput2 = 3; |
| 314 | const int kTensorCount = 4; |
| 315 | |
| 316 | // kInput1(0) --> +-------+ --> kOutput1(2) |
| 317 | // | WHILE | |
| 318 | // kInput2(1) --> +-------+ --> kOutput2(3) |
| 319 | |
| 320 | int first_new_tensor_index; |
| 321 | ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index), |
| 322 | kTfLiteOk); |
| 323 | ASSERT_EQ(first_new_tensor_index, 0); |
| 324 | ASSERT_EQ(subgraph->SetInputs({kInput1, kInput2}), kTfLiteOk); |
| 325 | ASSERT_EQ(subgraph->SetOutputs({kOutput1, kOutput2}), kTfLiteOk); |
| 326 | |
| 327 | SetupTensor(subgraph, kInput1, kTfLiteInt32); |
| 328 | SetupTensor(subgraph, kInput2, kTfLiteInt32); |
| 329 | SetupTensor(subgraph, kOutput1, kTfLiteInt32); |
| 330 | SetupTensor(subgraph, kOutput2, kTfLiteInt32); |
| 331 | |
| 332 | TfLiteWhileParams* params = |
| 333 | reinterpret_cast<TfLiteWhileParams*>(malloc(sizeof(TfLiteWhileParams))); |
| 334 | params->cond_subgraph_index = 1; |
| 335 | params->body_subgraph_index = 2; |
| 336 | |
| 337 | int node_index; |
| 338 | subgraph->AddNodeWithParameters({0, 1}, {2, 3}, {}, nullptr, 0, params, |
| 339 | ::tflite::ops::custom::Register_WHILE(), |
| 340 | &node_index); |
| 341 | } |
| 342 | |
| 343 | void SubgraphBuilder::CreateConstantInt32Tensor(Subgraph* subgraph, |
| 344 | int tensor_index, |
no test coverage detected