zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +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

13
python/tests/test_array.py
View File

@@ -1043,6 +1043,19 @@ class TestArray(mlx_tests.MLXTestCase):
a_mlx = mx.array(a_np)
self.assertTrue(np.array_equal(a_np[2:-1, 0], np.array(a_mlx[2:-1, 0])))
def test_indexing_grad(self):
x = mx.array([[1, 2], [3, 4]]).astype(mx.float32)
ind = mx.array([0, 1, 0]).astype(mx.float32)
def index_fn(x, ind):
return x[ind.astype(mx.int32)].sum()
grad_x, grad_ind = mx.grad(index_fn, argnums=(0, 1))(x, ind)
expected = mx.array([[2, 2], [1, 1]])
self.assertTrue(mx.array_equal(grad_x, expected))
self.assertTrue(mx.array_equal(grad_ind, mx.zeros(ind.shape)))
def test_setitem(self):
a = mx.array(0)
a[None] = 1