| 122 | } |
| 123 | |
| 124 | TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { |
| 125 | OpData* op_data = reinterpret_cast<OpData*>(node->user_data); |
| 126 | int num_inputs = node->inputs->size; |
| 127 | // The number of outputs should be the same as number of inputs. |
| 128 | TF_LITE_ENSURE_EQ(context, node->outputs->size, num_inputs); |
| 129 | |
| 130 | // Check subgraph indices and get subgraphs. |
| 131 | Subgraph* this_subgraph = reinterpret_cast<Subgraph*>(context->impl_); |
| 132 | auto* subgraphs = this_subgraph->GetSubgraphs(); |
| 133 | TF_LITE_ENSURE(context, op_data->cond_subgraph_index < subgraphs->size()); |
| 134 | TF_LITE_ENSURE(context, op_data->body_subgraph_index < subgraphs->size()); |
| 135 | |
| 136 | Subgraph* cond_subgraph = (*subgraphs)[op_data->cond_subgraph_index].get(); |
| 137 | Subgraph* body_subgraph = (*subgraphs)[op_data->body_subgraph_index].get(); |
| 138 | |
| 139 | // Check input & output count of the condition subgraph. |
| 140 | TF_LITE_ENSURE_EQ(context, cond_subgraph->inputs().size(), num_inputs); |
| 141 | TF_LITE_ENSURE_EQ(context, cond_subgraph->outputs().size(), 1); |
| 142 | |
| 143 | // Check input & output count of the body subgraph. |
| 144 | TF_LITE_ENSURE_EQ(context, body_subgraph->inputs().size(), num_inputs); |
| 145 | TF_LITE_ENSURE_EQ(context, body_subgraph->outputs().size(), num_inputs); |
| 146 | |
| 147 | // Prepare and check the condition subgraph. |
| 148 | TF_LITE_ENSURE_OK( |
| 149 | context, CopyTensorsShapeAndType( |
| 150 | context, this_subgraph, TfLiteIntArrayView(node->inputs), |
| 151 | cond_subgraph, cond_subgraph->inputs(), true)); |
| 152 | TF_LITE_ENSURE_OK(context, cond_subgraph->AllocateTensors()); |
| 153 | TfLiteTensor* cond_output = |
| 154 | cond_subgraph->tensor(cond_subgraph->outputs()[0]); |
| 155 | // TODO(ycling): Handle the case the cond subgraph has dynamic tensor outputs. |
| 156 | // This should rarely happens. In most cases the output is static with shape |
| 157 | // [1]. However theoretically intermediate tensors in the cond subgraph |
| 158 | // can be dynamic. |
| 159 | if (IsDynamicTensor(cond_output)) { |
| 160 | op_data->cond_has_dynamic_output_tensors = true; |
| 161 | } else { |
| 162 | TF_LITE_ENSURE_STATUS(CheckCondOutput(context, cond_output)); |
| 163 | } |
| 164 | |
| 165 | // Prepare and check the body subgraph. |
| 166 | TF_LITE_ENSURE_OK( |
| 167 | context, CopyTensorsShapeAndType( |
| 168 | context, this_subgraph, TfLiteIntArrayView(node->inputs), |
| 169 | body_subgraph, body_subgraph->inputs(), true)); |
| 170 | TF_LITE_ENSURE_OK(context, body_subgraph->AllocateTensors()); |
| 171 | if (body_subgraph->HasDynamicTensors()) { |
| 172 | op_data->body_has_dynamic_output_tensors = true; |
| 173 | } else { |
| 174 | for (int i = 0; i < num_inputs; ++i) { |
| 175 | TfLiteTensor* body_input = |
| 176 | body_subgraph->tensor(body_subgraph->inputs()[i]); |
| 177 | TfLiteTensor* body_output = |
| 178 | body_subgraph->tensor(body_subgraph->outputs()[i]); |
| 179 | TF_LITE_ENSURE_EQ(context, body_input->type, body_output->type); |
| 180 | |
| 181 | // TODO(ycling): Support dynamic sized body subgraph. |
nothing calls this directly
no test coverage detected