(
lam: Union[Tensor, float], size: Optional[Iterable[int]], seed: int, handle: int
)
| 191 | |
| 192 | |
| 193 | def _poisson( |
| 194 | lam: Union[Tensor, float], size: Optional[Iterable[int]], seed: int, handle: int |
| 195 | ) -> Tensor: |
| 196 | handle_cn = None if handle == 0 else _get_rng_handle_compnode(handle) |
| 197 | if not isinstance(lam, Tensor): |
| 198 | assert lam > 0, "Poisson is not defined when lam <= 0" |
| 199 | lam = Tensor(lam, dtype="float32", device=handle_cn) |
| 200 | if isinstance(size, int) and size != 0: |
| 201 | size = (size,) |
| 202 | assert ( |
| 203 | handle_cn is None or handle_cn == lam.device |
| 204 | ), "The lam ({}) must be the same device with handle ({})".format( |
| 205 | lam.device, handle_cn |
| 206 | ) |
| 207 | (lam,) = _broadcast_tensors_with_size([lam], size) |
| 208 | op = PoissonRNG(seed=seed, handle=handle) |
| 209 | (output,) = apply(op, lam) |
| 210 | return output |
| 211 | |
| 212 | |
| 213 | def _multinomial( |
no test coverage detected