Lines Matching refs:k
29 bool evalGeneric(const T* inputData, const Shape& inputShape, const int32_t k, T* valuesData, in evalGeneric() argument
41 std::nth_element(values.begin(), values.begin() + (rowSize - k), values.end()); in evalGeneric()
42 std::sort(values.begin() + (rowSize - k), values.end()); in evalGeneric()
44 for (int i = 0; i < k; ++i) { in evalGeneric()
56 bool prepare(const Shape& input, int32_t k, Shape* values, Shape* indices) { in prepare() argument
57 NN_CHECK(k > 0); in prepare()
58 NN_CHECK(k <= input.dimensions.back()); in prepare()
61 values->dimensions.back() = k; in prepare()
63 indices->dimensions.back() = k; in prepare()
67 bool eval(const void* inputData, const Shape& inputShape, const int32_t k, void* valuesData, in eval() argument
71 return evalGeneric(reinterpret_cast<const _Float16*>(inputData), inputShape, k, in eval()
76 return evalGeneric(reinterpret_cast<const float*>(inputData), inputShape, k, in eval()
81 return evalGeneric(reinterpret_cast<const int32_t*>(inputData), inputShape, k, in eval()
86 return evalGeneric(reinterpret_cast<const uint8_t*>(inputData), inputShape, k, in eval()