#NEWLINE# def test_take( ) : #NEWLINE# #INDENT# ( chunks , dsk ) = take ( 'y' , 'x' , [ ( 20 , 20 , 20 , 20 ) ] , [ 5 , 1 , 47 , 3 ] , axis = 0 ) #NEWLINE# expected = { #NEWLINE# #INDENT# ( 'y' , 0 ) : ( getitem , ( np . concatenate , ( list , [ ( getitem , ( 'x' , 0 ) , ( [ 1 , 3 , 5 ] , ) ) , ( getitem , ( 'x' , 2 ) , ( [ 7 ] , ) ) ] ) , 0 ) , ( [ 2 , 0 , 3 , 1 ] , ) ) , #NEWLINE# #UNINDENT# } #NEWLINE# assert ( dsk == expected ) #NEWLINE# assert ( chunks == ( ( 4 , ) , ) ) #NEWLINE# ( chunks , dsk ) = take ( 'y' , 'x' , [ ( 20 , 20 , 20 , 20 ) , ( 20 , 20 ) ] , [ 5 , 1 , 47 , 3 ] , axis = 0 ) #NEWLINE# expected = dict ( ( ( ( 'y' , 0 , j ) , ( getitem , ( np . concatenate , ( list , [ ( getitem , ( 'x' , 0 , j ) , ( [ 1 , 3 , 5 ] , slice ( None , None , None ) ) ) , ( getitem , ( 'x' , 2 , j ) , ( [ 7 ] , slice ( None , None , None ) ) ) ] ) , 0 ) , ( [ 2 , 0 , 3 , 1 ] , slice ( None , None , None ) ) ) ) for j in range ( 2 ) ) ) #NEWLINE# assert ( dsk == expected ) #NEWLINE# assert ( chunks == ( ( 4 , ) , ( 20 , 20 ) ) ) #NEWLINE# ( chunks , dsk ) = take ( 'y' , 'x' , [ ( 20 , 20 , 20 , 20 ) , ( 20 , 20 ) ] , [ 5 , 1 , 37 , 3 ] , axis = 1 ) #NEWLINE# expected = dict ( ( ( ( 'y' , i , 0 ) , ( getitem , ( np . concatenate , ( list , [ ( getitem , ( 'x' , i , 0 ) , ( slice ( None , None , None ) , [ 1 , 3 , 5 ] ) ) , ( getitem , ( 'x' , i , 1 ) , ( slice ( None , None , None ) , [ 17 ] ) ) ] ) , 1 ) , ( slice ( None , None , None ) , [ 2 , 0 , 3 , 1 ] ) ) ) for i in range ( 4 ) ) ) #NEWLINE# assert ( dsk == expected ) #NEWLINE# assert ( chunks == ( ( 20 , 20 , 20 , 20 ) , ( 4 , ) ) )