| 1122 | self.assertTrue(np.array_equal(a_np[2:-1, 0], np.array(a_mlx[2:-1, 0]))) |
| 1123 | |
| 1124 | def test_indexing_grad(self): |
| 1125 | x = mx.array([[1, 2], [3, 4]]).astype(mx.float32) |
| 1126 | ind = mx.array([0, 1, 0]).astype(mx.float32) |
| 1127 | |
| 1128 | def index_fn(x, ind): |
| 1129 | return x[ind.astype(mx.int32)].sum() |
| 1130 | |
| 1131 | grad_x, grad_ind = mx.grad(index_fn, argnums=(0, 1))(x, ind) |
| 1132 | expected = mx.array([[2, 2], [1, 1]]) |
| 1133 | |
| 1134 | self.assertTrue(mx.array_equal(grad_x, expected)) |
| 1135 | self.assertTrue(mx.array_equal(grad_ind, mx.zeros(ind.shape))) |
| 1136 | |
| 1137 | def test_setitem(self): |
| 1138 | a = mx.array(0) |