| 78 | } |
| 79 | |
| 80 | TEST(TestImperative, PoissonRNGBasic) { |
| 81 | REQUIRE_XPU(2); |
| 82 | for (auto&& cn : {CompNode::load("xpu0"), CompNode::load("xpu1")}) { |
| 83 | TensorShape shape{5, 3000}; |
| 84 | HostTensorND lam{cn, shape, dtype::Float32()}; |
| 85 | auto lam_ptr = lam.ptr<float>(); |
| 86 | for (int i = 0; i < 5 * 3000; ++i) |
| 87 | lam_ptr[i] = 2; |
| 88 | SmallVector<TensorPtr> inputs{Tensor::make(lam)}; |
| 89 | check_rng_with_input_basic<PoissonRNG>(cn, inputs, 123); |
| 90 | } |
| 91 | } |
| 92 | |
| 93 | TEST(TestImperative, BetaRNGBasic) { |
| 94 | REQUIRE_XPU(2); |