| 150 | } |
| 151 | |
| 152 | inline array complex_normal( |
| 153 | Shape shape, |
| 154 | const std::optional<array>& loc, |
| 155 | const std::optional<array>& scale, |
| 156 | const std::optional<array>& key, |
| 157 | StreamOrDevice s) { |
| 158 | auto stream = to_stream(s); |
| 159 | auto low = array(std::nextafter(-1.0f, 0.0f), float32); |
| 160 | auto high = array(1.0f, float32); |
| 161 | shape.push_back(2); |
| 162 | auto samples = |
| 163 | erfinv(uniform(low, high, shape, float32, key, stream), stream); |
| 164 | samples = squeeze(view(samples, complex64, stream), -1, stream); |
| 165 | if (scale.has_value()) { |
| 166 | samples = multiply(*scale, samples, stream); |
| 167 | } |
| 168 | if (loc.has_value()) { |
| 169 | samples = add(*loc, samples, stream); |
| 170 | } |
| 171 | return samples; |
| 172 | } |
| 173 | |
| 174 | array normal( |
| 175 | const Shape& shape, |