| 148 | } |
| 149 | |
| 150 | void SubgraphBuilder::BuildIfSubgraph(Subgraph* subgraph) { |
| 151 | const int kCondInput = 0; |
| 152 | const int kInput1 = 1; |
| 153 | const int kInput2 = 2; |
| 154 | const int kOutput = 3; |
| 155 | const int kTensorCount = 4; |
| 156 | |
| 157 | // kCondInput(0) --> +----+ |
| 158 | // kInput1(1) ----> | IF | --> kOutput(3) |
| 159 | // kInput2(2) ----> +----+ |
| 160 | |
| 161 | int first_new_tensor_index; |
| 162 | ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index), |
| 163 | kTfLiteOk); |
| 164 | ASSERT_EQ(first_new_tensor_index, 0); |
| 165 | ASSERT_EQ(subgraph->SetInputs({kCondInput, kInput1, kInput2}), kTfLiteOk); |
| 166 | ASSERT_EQ(subgraph->SetOutputs({kOutput}), kTfLiteOk); |
| 167 | |
| 168 | SetupTensor(subgraph, kCondInput, kTfLiteBool); |
| 169 | SetupTensor(subgraph, kInput1, kTfLiteInt32); |
| 170 | SetupTensor(subgraph, kInput2, kTfLiteInt32); |
| 171 | SetupTensor(subgraph, kOutput, kTfLiteInt32); |
| 172 | |
| 173 | TfLiteIfParams* params = |
| 174 | reinterpret_cast<TfLiteIfParams*>(malloc(sizeof(TfLiteIfParams))); |
| 175 | params->then_subgraph_index = 1; |
| 176 | params->else_subgraph_index = 2; |
| 177 | |
| 178 | int node_index; |
| 179 | subgraph->AddNodeWithParameters( |
| 180 | {kCondInput, kInput1, kInput2}, {kOutput}, {}, nullptr, 0, params, |
| 181 | ::tflite::ops::custom::Register_IF(), &node_index); |
| 182 | } |
| 183 | |
| 184 | void SubgraphBuilder::BuildLessEqualCondSubgraph(Subgraph* subgraph, int rhs) { |
| 185 | const int kInput1 = 0; |
no test coverage detected