Function run_multinomial_without_replacement

dnn/test/cuda/rng.cpp:153–205

template <typename T>
void run_multinomial_without_replacement(Handle* handle) {
    using ctype = typename DTypeTrait<T>::ctype;

    size_t num_groups = 2;
    size_t num_samples = 1;
    size_t len_probs = 4;
    bool replacement = false;
    size_t total_count = 10000;

    TensorLayout ly_probs{TensorShape{num_groups, len_probs}, T()};
    SyncedTensor<ctype> probs(handle, ly_probs);
    auto probs_ptr = probs.ptr_mutable_host();
    probs_ptr[0] = 1;
    probs_ptr[1] = 2;
    probs_ptr[2] = 3;
    probs_ptr[3] = 4;
    probs_ptr[4] = 0;
    probs_ptr[5] = 7;
    probs_ptr[6] = 2;
    probs_ptr[7] = 1;

    std::vector<float> norm_probs;
    for (size_t i = 0; i < 8; ++i) {
        norm_probs.push_back(probs_ptr[i] / 10);
    }

    auto opr = handle->create_operator<MultinomialRNG>();
    opr->param().num_samples = num_samples;
    opr->param().replacement = replacement;
    TensorLayout ly_out{TensorShape{num_groups, num_samples}, dtype::Int32()};
    SyncedTensor<dt_int32> out(handle, ly_out);
    Tensor<dt_byte> workspace(
            handle, {TensorShape{opr->get_workspace_in_bytes(ly_probs, ly_out)},
                     dtype::Byte()});
    std::vector<float> sample_probs(num_groups * len_probs, 0);
    for (size_t i = 0; i < total_count; ++i) {
        opr->exec(
                probs.tensornd_dev(), out.tensornd_dev(),
                {workspace.ptr(), workspace.layout().total_nr_elems()});

        auto ptr = out.ptr_mutable_host();
        sample_probs[ptr[0]] += 1;
        sample_probs[len_probs + ptr[1]] += 1;
    }

    for (size_t i = 0; i < num_groups * len_probs; ++i) {
        sample_probs[i] /= total_count * num_samples;
    }

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

template <typename T>
void run_beta(Handle* handle) {