| 109 | |
| 110 | template <typename T> |
| 111 | void run_multinomial(Handle* handle) { |
| 112 | using ctype = typename DTypeTrait<T>::ctype; |
| 113 | auto opr = handle->create_operator<MultinomialRNG>(); |
| 114 | |
| 115 | size_t num_groups = 2; |
| 116 | size_t num_samples = 10000; |
| 117 | size_t len_probs = 4; |
| 118 | bool replacement = true; |
| 119 | |
| 120 | TensorLayout ly_out{TensorShape{num_groups, num_samples}, dtype::Int32()}; |
| 121 | TensorLayout ly_probs{TensorShape{num_groups, len_probs}, T()}; |
| 122 | |
| 123 | Tensor<dt_int32> out(handle, ly_out); |
| 124 | Tensor<ctype> probs(handle, ly_probs); |
| 125 | |
| 126 | auto probs_ptr = probs.ptr(); |
| 127 | probs_ptr[0] = 0.1; |
| 128 | probs_ptr[1] = 0.2; |
| 129 | probs_ptr[2] = 0.3; |
| 130 | probs_ptr[3] = 0.4; |
| 131 | probs_ptr[4] = 0.0; |
| 132 | probs_ptr[5] = 0.7; |
| 133 | probs_ptr[6] = 0.2; |
| 134 | probs_ptr[7] = 0.1; |
| 135 | |
| 136 | opr->param().num_samples = num_samples; |
| 137 | opr->param().replacement = replacement; |
| 138 | opr->exec(probs.tensornd(), out.tensornd(), {}); |
| 139 | |
| 140 | auto ptr = out.ptr(); |
| 141 | |
| 142 | std::vector<float> sample_probs(num_groups * len_probs, 0); |
| 143 | for (size_t i = 0; i < num_groups; ++i) { |
| 144 | for (size_t j = 0; j < num_samples; ++j) { |
| 145 | sample_probs[i * len_probs + ptr[i * num_samples + j]] += 1; |
| 146 | } |
| 147 | } |
| 148 | for (size_t i = 0; i < num_groups * len_probs; ++i) { |
| 149 | sample_probs[i] /= num_samples; |
| 150 | } |
| 151 | |
| 152 | for (size_t i = 0; i < num_groups * len_probs; ++i) { |
| 153 | ASSERT_LE(std::abs(sample_probs[i] - probs_ptr[i]), 1e-2); |
| 154 | } |
| 155 | |
| 156 | std::vector<float> float_data_group0; |
| 157 | std::vector<float> float_data_group1; |
| 158 | for (size_t i = 0; i < num_samples; ++i) { |
| 159 | float_data_group0.push_back(static_cast<float>(ptr[i])); |
| 160 | } |
| 161 | for (size_t i = num_samples; i < 2 * num_samples; ++i) { |
| 162 | float_data_group1.push_back(static_cast<float>(ptr[i])); |
| 163 | } |
| 164 | float compare_mean_group0 = 0 * 0.1 + 1 * 0.2 + 2 * 0.3 + 3 * 0.4; |
| 165 | float compare_mean_group1 = 0 * 0.0 + 1 * 0.7 + 2 * 0.2 + 3 * 0.1; |
| 166 | float compare_var_group0 = (0 * 0.1 + 1 * 0.2 + 4 * 0.3 + 9 * 0.4) - |
| 167 | compare_mean_group0 * compare_mean_group0; |
| 168 | float compare_var_group1 = (0 * 0.0 + 1 * 0.7 + 4 * 0.2 + 9 * 0.1) - |