def test_unravel_index(): x1, y1 = rand_coord_2d((LARGE_X - 100), LARGE_X, 10, SMALL_Y) x2, y2 = rand_coord_2d((LARGE_X - 200), LARGE_X, 9, SMALL_Y) x3, y3 = rand_coord_2d((LARGE_X - 300), LARGE_X, 8, SMALL_Y) original_2d_indices = [[x1, x2, x3], [y1, y2, y3]] idx_numpy = np.ravel_multi_index(original_2d_indices, (LARGE_X, SMALL_Y)) indices_2d = mx.nd.unravel_index(mx.nd.array(idx_numpy, dtype=np.int64), shape=(LARGE_X, SMALL_Y)) assert (indices_2d.asnumpy() == np.array(original_2d_indices)).all()
def test_ravel_multi_index(): x1, y1 = rand_coord_2d((LARGE_X - 100), LARGE_X, 10, SMALL_Y) x2, y2 = rand_coord_2d((LARGE_X - 200), LARGE_X, 9, SMALL_Y) x3, y3 = rand_coord_2d((LARGE_X - 300), LARGE_X, 8, SMALL_Y) indices_2d = [[x1, x2, x3], [y1, y2, y3]] idx = mx.nd.ravel_multi_index(mx.nd.array(indices_2d, dtype=np.int64), shape=(LARGE_X, SMALL_Y)) idx_numpy = np.ravel_multi_index(indices_2d, (LARGE_X, SMALL_Y)) assert np.sum(1 for i in range(idx.size) if idx[i] == idx_numpy[i]) == 3