| 159 | |
| 160 | |
| 161 | def _beta( |
| 162 | alpha: Union[Tensor, float], |
| 163 | beta: Union[Tensor, float], |
| 164 | size: Optional[Iterable[int]], |
| 165 | seed: int, |
| 166 | handle: int, |
| 167 | ) -> Tensor: |
| 168 | handle_cn = None if handle == 0 else _get_rng_handle_compnode(handle) |
| 169 | if not isinstance(alpha, Tensor): |
| 170 | assert alpha > 0, "Beta is not defined when alpha <= 0" |
| 171 | alpha = Tensor(alpha, dtype="float32", device=handle_cn) |
| 172 | if not isinstance(beta, Tensor): |
| 173 | assert beta > 0, "Beta is not defined when beta <= 0" |
| 174 | beta = Tensor(beta, dtype="float32", device=handle_cn) |
| 175 | assert ( |
| 176 | handle_cn is None or handle_cn == alpha.device |
| 177 | ), "The alpha ({}) must be the same device with handle ({})".format( |
| 178 | alpha.device, handle_cn |
| 179 | ) |
| 180 | assert ( |
| 181 | handle_cn is None or handle_cn == beta.device |
| 182 | ), "The beta ({}) must be the same device with handle ({})".format( |
| 183 | beta.device, handle_cn |
| 184 | ) |
| 185 | if isinstance(size, int) and size != 0: |
| 186 | size = (size,) |
| 187 | alpha, beta = _broadcast_tensors_with_size([alpha, beta], size) |
| 188 | op = BetaRNG(seed=seed, handle=handle) |
| 189 | (output,) = apply(op, alpha, beta) |
| 190 | return output |
| 191 | |
| 192 | |
| 193 | def _poisson( |