def test_faiss_query(setup, X, Y, metric): faiss_index = build_faiss_index(X, 'Flat', None, metric=metric, random_state=0) d, i = faiss_query(faiss_index, Y, 5, nprobe=10) distance, indices = fetch(*execute(d, i)) nn = NearestNeighbors(metric=metric) nn.fit(x) expected_distance, expected_indices = nn.kneighbors(y, 5) np.testing.assert_array_equal(indices, expected_indices.fetch()) np.testing.assert_almost_equal(distance, expected_distance.fetch(), decimal=4) # test other index X2 = X.astype(np.float64) Y2 = y.astype(np.float64) faiss_index = build_faiss_index(X2, 'PCAR6,IVF8_HNSW32,SQ8', 10, random_state=0, return_index_type='object') d, i = faiss_query(faiss_index, Y2, 5, nprobe=10) # test execute only execute(d, i)
def test_pairwise_distances_topk_execution(setup): rs = np.random.RandomState(0) raw_x = rs.rand(20, 5) raw_y = rs.rand(21, 5) x = mt.tensor(raw_x, chunk_size=11) y = mt.tensor(raw_y, chunk_size=12) d, i = pairwise_distances_topk(x, y, 3, metric='euclidean', return_index=True) result = fetch(*execute(d, i)) nn = SkNearestNeighbors(n_neighbors=3, algorithm='brute', metric='euclidean') nn.fit(raw_y) expected = nn.kneighbors(raw_x, return_distance=True) np.testing.assert_almost_equal(result[0], expected[0]) np.testing.assert_array_equal(result[1], expected[1]) x = mt.tensor(raw_x, chunk_size=(11, 3)) d = pairwise_distances_topk(x, k=4, metric='euclidean', return_index=False) result = d.execute().fetch() nn = SkNearestNeighbors(n_neighbors=3, algorithm='brute', metric='euclidean') nn.fit(raw_x) expected = nn.kneighbors(return_distance=True)[0] np.testing.assert_almost_equal(result[:, 1:], expected) y = mt.tensor(raw_y, chunk_size=21) d, i = pairwise_distances_topk(x, y, 3, metric='cosine', return_index=True, working_memory='168') result = fetch(*execute(d, i)) nn = SkNearestNeighbors(n_neighbors=3, algorithm='brute', metric='cosine') nn.fit(raw_y) expected = nn.kneighbors(raw_x, return_distance=True) np.testing.assert_almost_equal(result[0], expected[0]) np.testing.assert_array_equal(result[1], expected[1]) d = pairwise_distances_topk(x, y, 3, metric='cosine', axis=0, return_index=False) result = d.execute().fetch() nn = SkNearestNeighbors(n_neighbors=3, algorithm='brute', metric='cosine') nn.fit(raw_x) expected = nn.kneighbors(raw_y, return_distance=True)[0] np.testing.assert_almost_equal(result, expected)
def test_no_default_session(): raw = np.random.RandomState(0).rand(10, 10) a = mt.tensor(raw, chunk_size=5) b = a + 1 with pytest.warns(Warning): execute(b, show_progress=False) np.testing.assert_array_equal(fetch(b), raw + 1) stop_server()
def _cancel_when_tile(session, cancelled): a = mt.tensor([1, 2, 3]) for i in range(20): a = SlowTileAdd(dtype=np.dtype(np.int64))(a, 1) execute(a, cancelled=cancelled) assert not a._executed_sessions del a ref_counts = session._get_ref_counts() assert len(ref_counts) == 0
def _cancel_when_execute(session, cancelled): def run(): time.sleep(200) rs = [mr.spawn(run) for _ in range(10)] execute(*rs, cancelled=cancelled) assert all(not r._executed_sessions for r in rs) del rs ref_counts = session._get_ref_counts() assert len(ref_counts) == 0
def test_modf_order_execution(setup): data1 = np.random.random((5, 9)) t = tensor(data1, chunk_size=3) o1, o2 = modf(t, order='F') res1, res2 = execute(o1, o2) expected1, expected2 = np.modf(data1, order='F') np.testing.assert_allclose(res1, expected1) assert res1.flags['F_CONTIGUOUS'] is True assert res1.flags['C_CONTIGUOUS'] is False np.testing.assert_allclose(res2, expected2) assert res2.flags['F_CONTIGUOUS'] is True assert res2.flags['C_CONTIGUOUS'] is False