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 ANDROID_PACKAGES_MODULES_NEURALNETWORKS_COMMON_TYPES_OPERATIONS_MULTINOMIAL_H
18 #define ANDROID_PACKAGES_MODULES_NEURALNETWORKS_COMMON_TYPES_OPERATIONS_MULTINOMIAL_H
19 
20 #include <algorithm>
21 #include <cmath>
22 #include <vector>
23 
24 #include "OperationsValidationUtils.h"
25 #include "nnapi/Types.h"
26 
27 namespace android {
28 namespace nn {
29 
30 struct RunTimeOperandInfo;
31 struct Shape;
32 
33 class Multinomial {
34    public:
35     Multinomial(const Operation& operation, RunTimeOperandInfo* operands);
36 
37     static bool Prepare(const Operation& operation, RunTimeOperandInfo* operands,
38                         Shape* outputShape);
39     bool Eval();
40 
41     static constexpr int kInputTensor = 0;
42     static constexpr int kSampleCountParam = 1;
43     static constexpr int kRandomSeedsTensor = 2;
44 
45     static constexpr int kOutputTensor = 0;
46 
47    private:
48     void EvalFloat32(const float* inputData);
49 
50     RunTimeOperandInfo* input_;
51     int sample_count_;
52     RunTimeOperandInfo* random_seeds_;
53 
54     RunTimeOperandInfo* output_;
55 };
56 
57 }  // namespace nn
58 }  // namespace android
59 
60 #endif  // ANDROID_PACKAGES_MODULES_NEURALNETWORKS_COMMON_TYPES_OPERATIONS_MULTINOMIAL_H
61