zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-11-01 08:38:12 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Real and Imag (#1490)

Browse Source 
* real and imag

* fix

* fix
This commit is contained in:
Awni Hannun
2024-10-15 16:23:15 -07:00
committed by GitHub
parent 2b8ace6a03
commit 3f86399922
21 changed files with 275 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
python/tests/test_autograd.py
View File

@@ -590,6 +590,21 @@ class TestAutograd(mlx_tests.MLXTestCase):
self.assertTrue(mx.allclose(out1[0], out2[0]))
self.assertTrue(mx.allclose(dout1[0] + 1, dout2[0]))
def test_complex_vjps(self):
def fun(x):
return (2.0 * mx.real(x)).sum()
x = mx.array([0.0 + 1j, 1.0 + 0.0j, 0.5 + 0.5j])
dfdx = mx.grad(fun)(x)
self.assertTrue(mx.allclose(dfdx, 2 * mx.ones_like(x)))
def fun(x):
return (2.0 * mx.imag(x)).sum()
x = mx.array([0.0 + 1j, 1.0 + 0.0j, 0.5 + 0.5j])
dfdx = mx.grad(fun)(x)
self.assertTrue(mx.allclose(dfdx, -2j * mx.ones_like(x)))
if __name__ == "__main__":
unittest.main()