| 7 | |
| 8 | template <typename In, typename Out> |
| 9 | static void multinomial_kernel(const In* input, |
| 10 | dim_t batch_size, |
| 11 | dim_t class_size, |
| 12 | dim_t sample_size, |
| 13 | Out* output) { |
| 14 | auto& generator = get_random_generator(); |
| 15 | |
| 16 | for (dim_t i = 0; i < batch_size; ++i) { |
| 17 | const In* input_data = input + i * class_size; |
| 18 | Out* output_data = output + i * sample_size; |
| 19 | |
| 20 | std::discrete_distribution<Out> distribution(input_data, input_data + class_size); |
| 21 | for (dim_t j = 0; j < sample_size; ++j) |
| 22 | output_data[j] = distribution(generator); |
| 23 | } |
| 24 | } |
| 25 | |
| 26 | template <Device D, typename T> |
| 27 | void Multinomial::compute(const StorageView& input, StorageView& output) const { |