| 1392 | a[:, None, 0] = mx.ones((1, 2)) |
| 1393 | |
| 1394 | def test_array_at(self): |
| 1395 | a = mx.array(1) |
| 1396 | with self.assertRaises(ValueError): |
| 1397 | a.at.add(1) |
| 1398 | |
| 1399 | a = a.at[None].add(1) |
| 1400 | self.assertEqual(a.item(), 2) |
| 1401 | |
| 1402 | a = mx.array([0, 1, 2]) |
| 1403 | a = a.at[1].add(2) |
| 1404 | self.assertEqual(a.tolist(), [0, 3, 2]) |
| 1405 | |
| 1406 | a = a.at[mx.array([0, 0, 0, 0])].add(1) |
| 1407 | self.assertEqual(a.tolist(), [4, 3, 2]) |
| 1408 | |
| 1409 | a = mx.zeros((10, 10)) |
| 1410 | a = a.at[0].add(mx.arange(10)) |
| 1411 | self.assertEqual(a[0].tolist(), list(range(10))) |
| 1412 | |
| 1413 | a = mx.zeros((10, 10)) |
| 1414 | index_x = mx.array([0, 2, 3, 7]) |
| 1415 | index_y = mx.array([3, 3, 1, 2]) |
| 1416 | u = mx.random.uniform(shape=(4,)) |
| 1417 | a = a.at[index_x, index_y].add(u) |
| 1418 | self.assertTrue(mx.allclose(a.sum(), u.sum())) |
| 1419 | self.assertEqualArray(a.sum(), u.sum(), atol=1e-6, rtol=1e-5) |
| 1420 | self.assertEqual(a[index_x, index_y].tolist(), u.tolist()) |
| 1421 | |
| 1422 | # Test all array.at ops |
| 1423 | a = mx.random.uniform(shape=(10, 5, 2)) |
| 1424 | idx_x = mx.array([0, 4]) |
| 1425 | update = mx.ones((2, 5)) |
| 1426 | a[idx_x, :, 0] = 0 |
| 1427 | a = a.at[idx_x, :, 0].add(update) |
| 1428 | self.assertEqualArray(a[idx_x, :, 0], update) |
| 1429 | a = a.at[idx_x, :, 0].subtract(update) |
| 1430 | self.assertEqualArray(a[idx_x, :, 0], mx.zeros_like(update)) |
| 1431 | a = a.at[idx_x, :, 0].add(2 * update) |
| 1432 | self.assertEqualArray(a[idx_x, :, 0], 2 * update) |
| 1433 | a = a.at[idx_x, :, 0].multiply(2 * update) |
| 1434 | self.assertEqualArray(a[idx_x, :, 0], 4 * update) |
| 1435 | a = a.at[idx_x, :, 0].divide(3 * update) |
| 1436 | self.assertEqualArray(a[idx_x, :, 0], (4 / 3) * update) |
| 1437 | a[idx_x, :, 0] = 5 |
| 1438 | update = mx.arange(10).reshape(2, 5) |
| 1439 | a = a.at[idx_x, :, 0].maximum(update) |
| 1440 | self.assertEqualArray(a[idx_x, :, 0], mx.maximum(a[idx_x, :, 0], update)) |
| 1441 | a[idx_x, :, 0] = 5 |
| 1442 | a = a.at[idx_x, :, 0].minimum(update) |
| 1443 | self.assertEqualArray(a[idx_x, :, 0], mx.minimum(a[idx_x, :, 0], update)) |
| 1444 | |
| 1445 | update = mx.array([1.0, 2.0])[None, None, None] |
| 1446 | src = mx.array([1.0, 2.0])[None, :] |
| 1447 | src = src.at[0:1].add(update) |
| 1448 | self.assertTrue(mx.array_equal(src, mx.array([[2.0, 4.0]]))) |
| 1449 | |
| 1450 | # Test all array.at ops with slice-only indices |
| 1451 | a = mx.random.uniform(shape=(10, 5, 2)) |