#include <stdint.h>
#include <algorithm>
#include <functional>
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
#include "tensorflow/lite/kernels/kernel_util.h"
namespace tflite {
namespace ops {
namespace builtin {
namespace unsorted_segment {
enum SegmentType { … };
static const int kInputDataTensor = …;
static const int kInputSegmentIdsTensor = …;
static const int kInputNumSegmentsTensor = …;
static const int kOutputTensor = …;
TfLiteStatus ResizeOutputTensor(TfLiteContext* context,
const TfLiteTensor* data,
const TfLiteTensor* segment_ids,
const TfLiteTensor* num_segments,
TfLiteTensor* output) { … }
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { … }
template <typename T>
struct SegmenMax { … };
template <typename T>
struct SegmenMin { … };
template <typename T>
struct SegmenProd { … };
template <typename T>
struct SegmenSum { … };
template <typename T>
TfLiteStatus EvalType(TfLiteContext* context, const RuntimeShape& input_shape,
const T* input_data,
const RuntimeShape& segment_ids_shape,
const int32_t* segment_ids_data,
const RuntimeShape& output_shape, T* output_data,
SegmentType segment_type) { … }
TfLiteStatus EvalGeneric(TfLiteContext* context, TfLiteNode* node,
SegmentType segment_type) { … }
TfLiteStatus EvalProd(TfLiteContext* context, TfLiteNode* node) { … }
TfLiteStatus EvalMax(TfLiteContext* context, TfLiteNode* node) { … }
TfLiteStatus EvalSum(TfLiteContext* context, TfLiteNode* node) { … }
TfLiteStatus EvalMin(TfLiteContext* context, TfLiteNode* node) { … }
}
TfLiteRegistration* Register_UNSORTED_SEGMENT_PROD() { … }
TfLiteRegistration* Register_UNSORTED_SEGMENT_MAX() { … }
TfLiteRegistration* Register_UNSORTED_SEGMENT_SUM() { … }
TfLiteRegistration* Register_UNSORTED_SEGMENT_MIN() { … }
}
}
}