zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-19 10:48:09 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix complex power on Metal (#1460)

Browse Source
This commit is contained in:
Awni Hannun
2024-10-06 19:58:30 -07:00
committed by GitHub
parent e4534dac17
commit 95d04805b3
2 changed files with 8 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
mlx/backend/metal/kernels/binary_ops.h
View File

@@ -217,7 +217,7 @@ struct Power {
template <>
complex64_t operator()(complex64_t x, complex64_t y) {
auto x_theta = metal::atan(x.imag / x.real);
auto x_theta = metal::atan2(x.imag, x.real);
auto x_ln_r = 0.5 * metal::log(x.real * x.real + x.imag * x.imag);
auto mag = metal::exp(y.real * x_ln_r - y.imag * x_theta);
auto phase = y.imag * x_ln_r + y.real * x_theta;