| 127 | |
| 128 | |
| 129 | def _gamma( |
| 130 | shape: Union[Tensor, float], |
| 131 | scale: Union[Tensor, float], |
| 132 | size: Optional[Iterable[int]], |
| 133 | seed: int, |
| 134 | handle: int, |
| 135 | ) -> Tensor: |
| 136 | handle_cn = None if handle == 0 else _get_rng_handle_compnode(handle) |
| 137 | if not isinstance(shape, Tensor): |
| 138 | assert shape > 0, "Gamma is not defined when shape <= 0" |
| 139 | shape = Tensor(shape, dtype="float32", device=handle_cn) |
| 140 | if not isinstance(scale, Tensor): |
| 141 | assert scale > 0, "Gamma is not defined when scale <= 0" |
| 142 | scale = Tensor(scale, dtype="float32", device=handle_cn) |
| 143 | assert ( |
| 144 | handle_cn is None or handle_cn == shape.device |
| 145 | ), "The shape ({}) must be the same device with handle ({})".format( |
| 146 | shape.device, handle_cn |
| 147 | ) |
| 148 | assert ( |
| 149 | handle_cn is None or handle_cn == scale.device |
| 150 | ), "The scale ({}) must be the same device with handle ({})".format( |
| 151 | scale.device, handle_cn |
| 152 | ) |
| 153 | if isinstance(size, int) and size != 0: |
| 154 | size = (size,) |
| 155 | shape, scale = _broadcast_tensors_with_size([shape, scale], size) |
| 156 | op = GammaRNG(seed=seed, handle=handle) |
| 157 | (output,) = apply(op, shape, scale) |
| 158 | return output |
| 159 | |
| 160 | |
| 161 | def _beta( |