You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
46 lines
2.1 KiB
46 lines
2.1 KiB
import pytest
|
|
import numpy as np
|
|
from dask.array.reshape import reshape_rechunk, expand_tuple, contract_tuple
|
|
|
|
|
|
@pytest.mark.parametrize('inshape,outshape,prechunks,inchunks,outchunks', [
|
|
((4,), (4,), ((2, 2),), ((2, 2),), ((2, 2),)),
|
|
((4,), (2, 2), ((2, 2),), ((2, 2),), ((1, 1), (2,))),
|
|
|
|
((4,), (4, 1), ((2, 2),), ((2, 2),), ((2, 2), (1,))),
|
|
((4,), (1, 4), ((2, 2),), ((2, 2),), ((1,), (2, 2))),
|
|
((1, 4), (4,), ((1,), (2, 2),), ((1,), (2, 2),), ((2, 2),)),
|
|
((4, 1), (4,), ((2, 2), (1,)), ((2, 2), (1,)), ((2, 2),)),
|
|
((4, 1, 4), (4, 4), ((2, 2), (1,), (2, 2),), ((2, 2), (1,), (2, 2),), ((2, 2), (2, 2))),
|
|
((4, 4), (4, 1, 4), ((2, 2), (2, 2),), ((2, 2), (2, 2),), ((2, 2), (1,), (2, 2))),
|
|
|
|
((2, 2), (4,), ((2,), (2,)), ((2,), (2,)), ((4,),)),
|
|
((2, 2), (4,), ((1, 1), (2,)), ((1, 1,), (2,)), ((2, 2),)),
|
|
((2, 2), (4,), ((2,), (1, 1)), ((1, 1,), (2,)), ((2, 2),)),
|
|
|
|
((64,), (4, 4, 4), ((8, 8, 8, 8, 8, 8, 8, 8),), ((16, 16, 16, 16),), ((1, 1, 1, 1), (4,), (4,))),
|
|
((64,), (4, 4, 4), ((32, 32),), ((32, 32),), ((2, 2), (4,), (4,))),
|
|
((64,), (4, 4, 4), ((16, 48),), ((16, 48),), ((1, 3), (4,), (4,))),
|
|
((64,), (4, 4, 4), ((20, 44),), ((16, 48),), ((1, 3), (4,), (4,))),
|
|
((64, 4), (8, 8, 4), ((16, 16, 16, 16), (2, 2)), ((16, 16, 16, 16), (2, 2)), ((2, 2, 2, 2), (8,), (2, 2))),
|
|
])
|
|
def test_reshape_rechunk(inshape, outshape, prechunks, inchunks, outchunks):
|
|
result_in, result_out = reshape_rechunk(inshape, outshape, prechunks)
|
|
assert result_in == inchunks
|
|
assert result_out == outchunks
|
|
assert np.prod(list(map(len, result_in))) == np.prod(list(map(len, result_out)))
|
|
|
|
|
|
def test_expand_tuple():
|
|
assert expand_tuple((2, 4), 2) == (1, 1, 2, 2)
|
|
assert expand_tuple((2, 4), 3) == (1, 1, 1, 1, 2)
|
|
assert expand_tuple((3, 4), 2) == (1, 2, 2, 2)
|
|
assert expand_tuple((7, 4), 3) == (2, 2, 3, 1, 1, 2)
|
|
|
|
|
|
def test_contract_tuple():
|
|
assert contract_tuple((1, 1, 2, 3, 1), 2) == (2, 2, 2, 2)
|
|
assert contract_tuple((1, 1, 2, 5, 1), 2) == (2, 2, 4, 2)
|
|
assert contract_tuple((2, 4), 2) == (2, 4)
|
|
assert contract_tuple((2, 4), 3) == (6,)
|