MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / SelectPrepare

Function SelectPrepare

tensorflow/lite/kernels/select.cc:32–62  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

30constexpr int kOutputTensor = 0;
31
32TfLiteStatus SelectPrepare(TfLiteContext* context, TfLiteNode* node) {
33 TF_LITE_ENSURE_EQ(context, NumInputs(node), 3);
34 TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
35
36 const TfLiteTensor* input_condition =
37 GetInput(context, node, kInputTensorCondition);
38 const TfLiteTensor* input_x = GetInput(context, node, kInputTensorX);
39 const TfLiteTensor* input_y = GetInput(context, node, kInputTensorY);
40 TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
41
42 // Input must be bool.
43 TF_LITE_ENSURE(context, input_condition->type == kTfLiteBool);
44
45 // Input tensors must have the same type and size
46 TF_LITE_ENSURE_EQ(context, input_x->type, input_y->type);
47 TF_LITE_ENSURE(context, HaveSameShapes(input_x, input_y));
48 output->type = input_x->type;
49
50 // Either the same shape, or input_condition must be Rank 1 and match over the
51 // first dimension.
52 bool same_shape = HaveSameShapes(input_condition, input_x);
53 if (!same_shape && NumDimensions(input_condition) == 1) {
54 same_shape =
55 SizeOfDimension(input_condition, 0) == SizeOfDimension(input_x, 0);
56 }
57
58 TF_LITE_ENSURE(context, same_shape);
59
60 TfLiteIntArray* output_size = TfLiteIntArrayCopy(input_x->dims);
61 return context->ResizeTensor(context, output, output_size);
62}
63
64TfLiteStatus SelectEval(TfLiteContext* context, TfLiteNode* node) {
65 const TfLiteTensor* input_condition =

Callers

nothing calls this directly

Calls 9

NumInputsFunction · 0.85
GetInputFunction · 0.85
GetOutputFunction · 0.85
HaveSameShapesFunction · 0.85
NumDimensionsFunction · 0.85
SizeOfDimensionFunction · 0.85
TfLiteIntArrayCopyFunction · 0.85
ResizeTensorMethod · 0.80
NumOutputsFunction · 0.70

Tested by

no test coverage detected