(
rate: Union[Tensor, float], size: Optional[Iterable[int]], seed: int, handle: int
)
| 265 | |
| 266 | |
| 267 | def _exponential( |
| 268 | rate: Union[Tensor, float], size: Optional[Iterable[int]], seed: int, handle: int |
| 269 | ) -> Tensor: |
| 270 | handle_cn = None if handle == 0 else _get_rng_handle_compnode(handle) |
| 271 | if not isinstance(rate, Tensor): |
| 272 | assert rate > 0, "Exponential is not defined when rate <= 0" |
| 273 | rate = Tensor(rate, dtype="float32", device=handle_cn) |
| 274 | if isinstance(size, int) and size != 0: |
| 275 | size = (size,) |
| 276 | assert ( |
| 277 | handle_cn is None or handle_cn == rate.device |
| 278 | ), "The rate ({}) must be the same device with handle ({})".format( |
| 279 | rate.device, handle_cn |
| 280 | ) |
| 281 | (rate,) = _broadcast_tensors_with_size([rate], size) |
| 282 | op = ExponentialRNG(seed=seed, handle=handle) |
| 283 | (output,) = apply(op, rate) |
| 284 | return output |
| 285 | |
| 286 | |
| 287 | class RNG: |
no test coverage detected