MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / fill_multinomial_without_replacement

Function fill_multinomial_without_replacement

dnn/src/naive/rng/opr_impl.cpp:125–151  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

123
124template <typename U>
125void fill_multinomial_without_replacement(
126 Xoroshiro128plus* rng, U* probs, dt_int32* dst, size_t num_groups,
127 size_t num_samples, size_t len_probs) {
128 std::vector<std::vector<U>> data(num_groups, std::vector<U>(len_probs, U(0)));
129 for (size_t i = 0; i < num_groups; ++i) {
130 for (size_t j = 0; j < len_probs; ++j) {
131 data[i][j] = log(uniform_sample<U>(rng)) / probs[i * len_probs + j];
132 }
133 }
134
135 std::vector<std::vector<dt_int32>> index(
136 num_groups, std::vector<dt_int32>(len_probs, 0));
137 for (size_t i = 0; i < num_groups; ++i) {
138 for (size_t j = 0; j < len_probs; ++j) {
139 index[i][j] = j;
140 }
141 }
142
143 for (size_t i = 0; i < num_groups; ++i) {
144 std::sort(index[i].begin(), index[i].end(), [&](size_t idx1, size_t idx2) {
145 return data[i][idx1] > data[i][idx2];
146 });
147 std::copy(
148 index[i].begin(), index[i].begin() + num_samples,
149 dst + i * num_samples);
150 }
151}
152
153template <typename U>
154void fill_multinomial(

Callers 1

fill_multinomialFunction · 0.85

Calls 6

sortFunction · 0.85
UFunction · 0.70
logFunction · 0.50
copyFunction · 0.50
beginMethod · 0.45
endMethod · 0.45

Tested by

no test coverage detected