1 /* 2 * Copyright (C) 2018 The Android Open Source Project 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 #ifndef FRAMEWORKS_ML_NN_MULTINOMIAL_H 18 #define FRAMEWORKS_ML_NN_MULTINOMIAL_H 19 20 #include "HalOperation.h" 21 22 #include "tensorflow/lite/kernels/internal/tensor_utils.h" 23 24 #include <algorithm> 25 #include <cmath> 26 27 namespace android { 28 namespace nn { 29 30 struct RunTimeOperandInfo; 31 struct Shape; 32 33 class Multinomial { 34 public: 35 Multinomial(const android::hardware::neuralnetworks::V1_2::Operation& operation, 36 std::vector<RunTimeOperandInfo>& operands); 37 38 static bool Prepare(const hardware::neuralnetworks::V1_2::Operation& operation, 39 std::vector<RunTimeOperandInfo>& operands, Shape* outputShape); 40 bool Eval(); 41 42 static constexpr int kInputTensor = 0; 43 static constexpr int kSampleCountParam = 1; 44 static constexpr int kRandomSeedsTensor = 2; 45 46 static constexpr int kOutputTensor = 0; 47 48 private: 49 void EvalFloat32(const float* inputData); 50 51 RunTimeOperandInfo* input_; 52 int sample_count_; 53 RunTimeOperandInfo* random_seeds_; 54 55 RunTimeOperandInfo* output_; 56 }; 57 58 } // namespace nn 59 } // namespace android 60 61 #endif // FRAMEWORKS_ML_NN_MULTINOMIAL_H 62
