| 78 | namespace { |
| 79 | |
| 80 | TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { |
| 81 | TF_LITE_ENSURE_EQ(context, NumInputs(node), 5); |
| 82 | TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); |
| 83 | |
| 84 | const TfLiteTensor* ids = GetInput(context, node, 0); |
| 85 | TF_LITE_ENSURE_EQ(context, NumDimensions(ids), 1); |
| 86 | TF_LITE_ENSURE_EQ(context, ids->type, kTfLiteInt32); |
| 87 | |
| 88 | const TfLiteTensor* indices = GetInput(context, node, 1); |
| 89 | TF_LITE_ENSURE_EQ(context, NumDimensions(indices), 2); |
| 90 | TF_LITE_ENSURE_EQ(context, indices->type, kTfLiteInt32); |
| 91 | |
| 92 | const TfLiteTensor* shape = GetInput(context, node, 2); |
| 93 | TF_LITE_ENSURE_EQ(context, NumDimensions(shape), 1); |
| 94 | TF_LITE_ENSURE_EQ(context, shape->type, kTfLiteInt32); |
| 95 | |
| 96 | const TfLiteTensor* weights = GetInput(context, node, 3); |
| 97 | TF_LITE_ENSURE_EQ(context, NumDimensions(weights), 1); |
| 98 | TF_LITE_ENSURE_EQ(context, weights->type, kTfLiteFloat32); |
| 99 | |
| 100 | TF_LITE_ENSURE_EQ(context, SizeOfDimension(indices, 0), |
| 101 | SizeOfDimension(ids, 0)); |
| 102 | TF_LITE_ENSURE_EQ(context, SizeOfDimension(indices, 0), |
| 103 | SizeOfDimension(weights, 0)); |
| 104 | |
| 105 | const TfLiteTensor* value = GetInput(context, node, 4); |
| 106 | TF_LITE_ENSURE(context, NumDimensions(value) >= 2); |
| 107 | |
| 108 | // Mark the output as a dynamic tensor. |
| 109 | TfLiteTensor* output = GetOutput(context, node, 0); |
| 110 | TF_LITE_ENSURE_EQ(context, output->type, kTfLiteFloat32); |
| 111 | output->allocation_type = kTfLiteDynamic; |
| 112 | |
| 113 | return kTfLiteOk; |
| 114 | } |
| 115 | |
| 116 | void FinalizeAggregation(TfLiteCombinerType combiner, int num_elements, |
| 117 | float current_total_weight, |
nothing calls this directly
no test coverage detected