def test_topk():
x = np.array([5, 2, 1, 6])
d = da.from_array(x, chunks=2)
e = da.topk(2, d)
assert e.chunks == ((2,),)
assert eq(e, np.sort(x)[-1:-3:-1])
assert same_keys(da.topk(2, d), e)
def test_topk():
x = np.array([5, 2, 1, 6])
d = da.from_array(x, chunks=2)
e = da.topk(2, d)
assert e.chunks == ((2, ), )
assert_eq(e, np.sort(x)[-1:-3:-1])
assert same_keys(da.topk(2, d), e)
def test_topk_argtopk2():
a = da.random.random((10, 20, 30), chunks=(4, 8, 8))
# Support for deprecated API for topk
with pytest.warns(UserWarning):
assert_eq(da.topk(a, 5), da.topk(5, a))
# As Array methods
assert_eq(a.topk(5, axis=1, split_every=2), da.topk(a, 5, axis=1, split_every=2))
assert_eq(a.argtopk(5, axis=1, split_every=2), da.argtopk(a, 5, axis=1, split_every=2))
def test_topk_k_bigger_than_chunk():
x = np.array([5, 2, 1, 6])
d = da.from_array(x, chunks=2)
e = da.topk(3, d)
assert e.chunks == ((3,),)
assert eq(e, np.array([6, 5, 2]))
def test_topk_k_bigger_than_chunk():
x = np.array([5, 2, 1, 6])
d = da.from_array(x, chunks=2)
e = da.topk(3, d)
assert e.chunks == ((3, ), )
assert_eq(e, np.array([6, 5, 2]))
def test_topk_argtopk3():
a = da.random.random((10, 20, 30), chunks=(4, 8, 8))
# As Array methods
assert_eq(a.topk(5, axis=1, split_every=2),
da.topk(a, 5, axis=1, split_every=2))
assert_eq(a.argtopk(5, axis=1, split_every=2),
da.argtopk(a, 5, axis=1, split_every=2))
def test_topk_argtopk3():
a = da.random.random((10, 20, 30), chunks=(4, 8, 8))
# As Array methods
assert_eq(a.topk(5, axis=1, split_every=2),
da.topk(a, 5, axis=1, split_every=2))
assert_eq(a.argtopk(5, axis=1, split_every=2),
da.argtopk(a, 5, axis=1, split_every=2))
