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
