Function run_multinomial

dnn/test/cuda/rng.cpp:70–150

template <typename T>
void run_multinomial(Handle* handle) {
    using ctype = typename DTypeTrait<T>::ctype;
    auto opr = handle->create_operator<MultinomialRNG>();

    size_t num_groups = 2;
    size_t num_samples = 10000;
    size_t len_probs = 4;
    bool replacement = true;

    TensorLayout ly_out{TensorShape{num_groups, num_samples}, dtype::Int32()};
    TensorLayout ly_probs{TensorShape{num_groups, len_probs}, T()};
    SyncedTensor<dt_int32> out(handle, ly_out);
    SyncedTensor<ctype> probs(handle, ly_probs);

    auto probs_ptr = probs.ptr_mutable_host();
    probs_ptr[0] = 0.1;
    probs_ptr[1] = 0.2;
    probs_ptr[2] = 0.3;
    probs_ptr[3] = 0.4;
    probs_ptr[4] = 0.0;
    probs_ptr[5] = 0.7;
    probs_ptr[6] = 0.2;
    probs_ptr[7] = 0.1;

    Tensor<dt_byte> workspace(
            handle, {TensorShape{opr->get_workspace_in_bytes(ly_probs, ly_out)},
                     dtype::Byte()});

    opr->param().num_samples = num_samples;
    opr->param().replacement = replacement;
    opr->exec(
            probs.tensornd_dev(), out.tensornd_dev(),
            {workspace.ptr(), workspace.layout().total_nr_elems()});

    auto ptr = out.ptr_mutable_host();

    std::vector<float> sample_probs(num_groups * len_probs, 0);
    for (size_t i = 0; i < num_groups; ++i) {
        for (size_t j = 0; j < num_samples; ++j) {
            sample_probs[i * len_probs + ptr[i * num_samples + j]] += 1;
        }
    }
    for (size_t i = 0; i < num_groups * len_probs; ++i) {
        sample_probs[i] /= num_samples;
    }

    for (size_t i = 0; i < num_groups * len_probs; ++i) {
        ASSERT_LE(std::abs(sample_probs[i] - probs_ptr[i]), 1e-2);
    }

    std::vector<float> float_data_group0;
    std::vector<float> float_data_group1;
    for (size_t i = 0; i < num_samples; ++i) {
        float_data_group0.push_back(static_cast<float>(ptr[i]));
    }
    for (size_t i = num_samples; i < 2 * num_samples; ++i) {
        float_data_group1.push_back(static_cast<float>(ptr[i]));
    }