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

add zero vjps for bitwise ops and gather w.r.t. index (#1256)

Browse Source
This commit is contained in:
Alex Barron
2024-07-07 21:34:59 -07:00
committed by GitHub
parent 20bb301195
commit bdb36c9a63
4 changed files with 62 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
python/tests/test_ops.py
View File

@@ -2373,6 +2373,21 @@ class TestOps(mlx_tests.MLXTestCase):
self.assertEqual(c.shape, (3, 2, 5))
self.assertTrue(mx.array_equal(c, mx.ones((3, 2, 5), dtype=mx.bool_)))
def test_bitwise_grad(self):
a = np.random.randint(0, 10, size=(4, 3))
b = np.random.randint(0, 10, size=(4, 3))
cotangent = np.random.randint(0, 10, size=(4, 3))
a = mx.array(a)
b = mx.array(b)
cotangent = mx.array(cotangent)
def bitwise(a, b):
return a.astype(mx.int32) & b.astype(mx.int32)
_, vjps = mx.vjp(bitwise, [a, b], [cotangent])
for vjp in vjps:
self.assertFalse(np.any(np.array(vjp)))
def test_conjugate(self):
shape = (3, 5, 7)
a = np.random.normal(size=shape) + 1j * np.random.normal(size=shape)