L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Conjugate VJP and JVP support (#3386)

Browse Source
This commit is contained in:
Cameron Churchwell
2026-04-09 17:04:46 -05:00
committed by GitHub
parent 8332e228e4
commit d6d9b24801
3 changed files with 26 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tests/test_ops.py
+10
View File
@@ -2888,6 +2888,16 @@ class TestOps(mlx_tests.MLXTestCase):
out_np = a.conj()
self.assertTrue(np.array_equal(np.array(out_mlx), out_np))
b = np.random.normal(size=shape) + 1j * np.random.normal(size=shape)
b = b.astype(np.complex64)
_, vjps = mx.vjp(mx.conj, [mx.array(a)], [mx.array(b)])
self.assertTrue(np.array_equal(np.array(vjps[0]), b.conj()))
out_mlx, jvps = mx.jvp(mx.conj, [mx.array(a)], [mx.array(b)])
self.assertTrue(np.array_equal(np.array(out_mlx[0]), a.conj()))
self.assertTrue(np.array_equal(np.array(jvps[0]), b.conj()))
def test_view(self):
# Check scalar
out = mx.array(1, mx.int8).view(mx.uint8).item()