<s> import pytest <EOL> pytest . importorskip ( '<STR_LIT>' ) <EOL> from dask . array . optimization import ( getitem , rewrite_rules , optimize , <EOL> remove_full_slices , fuse_slice ) <EOL> from dask . utils import raises <EOL> from dask . array . core import getarray <EOL> def test_fuse_getitem ( ) : <EOL> pairs = [ ( ( getarray , ( getarray , '<STR_LIT:x>' , slice ( <NUM_LIT:1000> , <NUM_LIT> ) ) , slice ( <NUM_LIT:15> , <NUM_LIT:20> ) ) , <EOL> ( getarray , '<STR_LIT:x>' , slice ( <NUM_LIT> , <NUM_LIT> ) ) ) , <EOL> ( ( getitem , ( getarray , '<STR_LIT:x>' , ( slice ( <NUM_LIT:1000> , <NUM_LIT> ) , slice ( <NUM_LIT:100> , <NUM_LIT:200> ) ) ) , <EOL> ( slice ( <NUM_LIT:15> , <NUM_LIT:20> ) , slice ( <NUM_LIT:50> , <NUM_LIT> ) ) ) , <EOL> ( getarray , '<STR_LIT:x>' , ( slice ( <NUM_LIT> , <NUM_LIT> ) , slice ( <NUM_LIT> , <NUM_LIT> ) ) ) ) , <EOL> ( ( getarray , ( getarray , '<STR_LIT:x>' , slice ( <NUM_LIT:1000> , <NUM_LIT> ) ) , <NUM_LIT:10> ) , <EOL> ( getarray , '<STR_LIT:x>' , <NUM_LIT> ) ) , <EOL> ( ( getitem , ( getarray , '<STR_LIT:x>' , ( slice ( <NUM_LIT:1000> , <NUM_LIT> ) , <NUM_LIT:10> ) ) , <EOL> ( slice ( <NUM_LIT:15> , <NUM_LIT:20> ) , ) ) , <EOL> ( getarray , '<STR_LIT:x>' , ( slice ( <NUM_LIT> , <NUM_LIT> ) , <NUM_LIT:10> ) ) ) , <EOL> ( ( getarray , ( getarray , '<STR_LIT:x>' , ( <NUM_LIT:10> , slice ( <NUM_LIT:1000> , <NUM_LIT> ) ) ) , <EOL> ( slice ( <NUM_LIT:15> , <NUM_LIT:20> ) , ) ) , <EOL> ( getarray , '<STR_LIT:x>' , ( <NUM_LIT:10> , slice ( <NUM_LIT> , <NUM_LIT> ) ) ) ) , <EOL> ( ( getarray , ( getarray , '<STR_LIT:x>' , ( slice ( <NUM_LIT:1000> , <NUM_LIT> ) , slice ( <NUM_LIT:100> , <NUM_LIT:200> ) ) ) , <EOL> ( slice ( None , None ) , slice ( <NUM_LIT:50> , <NUM_LIT> ) ) ) , <EOL> ( getarray , '<STR_LIT:x>' , ( slice ( <NUM_LIT:1000> , <NUM_LIT> ) , slice ( <NUM_LIT> , <NUM_LIT> ) ) ) ) , <EOL> ( ( getarray , ( getarray , '<STR_LIT:x>' , ( None , slice ( None , None ) ) ) , <EOL> ( slice ( None , None ) , <NUM_LIT:5> ) ) , <EOL> ( getarray , '<STR_LIT:x>' , ( None , <NUM_LIT:5> ) ) ) , <EOL> ( ( getarray , ( getarray , '<STR_LIT:x>' , ( slice ( <NUM_LIT:1000> , <NUM_LIT> ) , slice ( <NUM_LIT:10> , <NUM_LIT:20> ) ) ) , <EOL> ( slice ( <NUM_LIT:5> , <NUM_LIT:10> ) , ) ) , <EOL> ( getarray , '<STR_LIT:x>' , ( slice ( <NUM_LIT> , <NUM_LIT> ) , slice ( <NUM_LIT:10> , <NUM_LIT:20> ) ) ) ) , <EOL> ( ( getitem , ( getitem , '<STR_LIT:x>' , ( slice ( <NUM_LIT:1000> , <NUM_LIT> ) , ) ) , <EOL> ( slice ( <NUM_LIT:5> , <NUM_LIT:10> ) , slice ( <NUM_LIT:10> , <NUM_LIT:20> ) ) ) , <EOL> ( getitem , '<STR_LIT:x>' , ( slice ( <NUM_LIT> , <NUM_LIT> ) , slice ( <NUM_LIT:10> , <NUM_LIT:20> ) ) ) ) <EOL> ] <EOL> for inp , expected in pairs : <EOL> result = rewrite_rules . rewrite ( inp ) <EOL> assert result == expected <EOL> def test_optimize_with_getitem_fusion ( ) : <EOL> dsk = { '<STR_LIT:a>' : '<STR_LIT>' , <EOL> '<STR_LIT:b>' : ( getarray , '<STR_LIT:a>' , ( slice ( <NUM_LIT:10> , <NUM_LIT:20> ) , slice ( <NUM_LIT:100> , <NUM_LIT:200> ) ) ) , <EOL> '<STR_LIT:c>' : ( getarray , '<STR_LIT:b>' , ( <NUM_LIT:5> , slice ( <NUM_LIT:50> , <NUM_LIT> ) ) ) } <EOL> result = optimize ( dsk , [ '<STR_LIT:c>' ] ) <EOL> expected = { '<STR_LIT:c>' : ( getarray , '<STR_LIT>' , ( <NUM_LIT:15> , slice ( <NUM_LIT> , <NUM_LIT> ) ) ) } <EOL> assert result == expected <EOL> def test_optimize_slicing ( ) : <EOL> dsk = { '<STR_LIT:a>' : ( range , <NUM_LIT:10> ) , <EOL> '<STR_LIT:b>' : ( getarray , '<STR_LIT:a>' , ( slice ( None , None , None ) , ) ) , <EOL> '<STR_LIT:c>' : ( getarray , '<STR_LIT:b>' , ( slice ( None , None , None ) , ) ) , <EOL> '<STR_LIT:d>' : ( getarray , '<STR_LIT:c>' , ( slice ( None , <NUM_LIT:5> , None ) , ) ) , <EOL> '<STR_LIT:e>' : ( getarray , '<STR_LIT:d>' , ( slice ( None , None , None ) , ) ) } <EOL> expected = { '<STR_LIT:a>' : ( range , <NUM_LIT:10> ) , <EOL> '<STR_LIT:e>' : ( getarray , '<STR_LIT:a>' , ( slice ( None , <NUM_LIT:5> , None ) , ) ) } <EOL> assert remove_full_slices ( dsk , [ ] ) == expected <EOL> expected = { '<STR_LIT:a>' : ( range , <NUM_LIT:10> ) , <EOL> '<STR_LIT:c>' : ( getarray , '<STR_LIT:a>' , ( slice ( None , None , None ) , ) ) , <EOL> '<STR_LIT:d>' : ( getarray , '<STR_LIT:c>' , ( slice ( None , <NUM_LIT:5> , None ) , ) ) , <EOL> '<STR_LIT:e>' : ( getarray , '<STR_LIT:d>' , ( slice ( None , None , None ) , ) ) } <EOL> assert remove_full_slices ( dsk , [ '<STR_LIT:c>' , '<STR_LIT:d>' , '<STR_LIT:e>' ] ) == expected <EOL> def test_fuse_slice ( ) : <EOL> assert fuse_slice ( slice ( <NUM_LIT:10> , <NUM_LIT:15> ) , slice ( <NUM_LIT:0> , <NUM_LIT:5> , <NUM_LIT:2> ) ) == slice ( <NUM_LIT:10> , <NUM_LIT:15> , <NUM_LIT:2> ) <EOL> assert fuse_slice ( ( slice ( <NUM_LIT:100> , <NUM_LIT:200> ) , ) , ( None , slice ( <NUM_LIT:10> , <NUM_LIT:20> ) ) ) == ( None , slice ( <NUM_LIT> , <NUM_LIT> ) ) <EOL> assert fuse_slice ( ( slice ( <NUM_LIT:100> , <NUM_LIT:200> ) , ) , ( slice ( <NUM_LIT:10> , <NUM_LIT:20> ) , None ) ) == ( slice ( <NUM_LIT> , <NUM_LIT> ) , None ) <EOL> assert fuse_slice ( ( <NUM_LIT:1> , ) , ( None , ) ) == ( <NUM_LIT:1> , None ) <EOL> assert fuse_slice ( ( <NUM_LIT:1> , slice ( <NUM_LIT:10> , <NUM_LIT:20> ) ) , ( None , None , <NUM_LIT:3> , None ) ) == ( <NUM_LIT:1> , None , None , <NUM_LIT> , None ) <EOL> assert ( raises ( NotImplementedError , <EOL> lambda : fuse_slice ( slice ( <NUM_LIT:10> , <NUM_LIT:15> , <NUM_LIT:2> ) , - <NUM_LIT:1> ) ) or <EOL> fuse_slice ( slice ( <NUM_LIT:10> , <NUM_LIT:15> , <NUM_LIT:2> ) , - <NUM_LIT:1> ) == <NUM_LIT> ) <EOL> def test_fuse_slice_with_lists ( ) : <EOL> assert fuse_slice ( slice ( <NUM_LIT:10> , <NUM_LIT:20> , <NUM_LIT:2> ) , [ <NUM_LIT:1> , <NUM_LIT:2> , <NUM_LIT:3> ] ) == [ <NUM_LIT:12> , <NUM_LIT> , <NUM_LIT:16> ] <EOL> assert fuse_slice ( [ <NUM_LIT:10> , <NUM_LIT:20> , <NUM_LIT:30> , <NUM_LIT> , <NUM_LIT:50> ] , [ <NUM_LIT:3> , <NUM_LIT:1> , <NUM_LIT:2> ] ) == [ <NUM_LIT> , <NUM_LIT:20> , <NUM_LIT:30> ] <EOL> assert fuse_slice ( [ <NUM_LIT:10> , <NUM_LIT:20> , <NUM_LIT:30> , <NUM_LIT> , <NUM_LIT:50> ] , <NUM_LIT:3> ) == <NUM_LIT> <EOL> assert fuse_slice ( [ <NUM_LIT:10> , <NUM_LIT:20> , <NUM_LIT:30> , <NUM_LIT> , <NUM_LIT:50> ] , - <NUM_LIT:1> ) == <NUM_LIT:50> <EOL> assert fuse_slice ( [ <NUM_LIT:10> , <NUM_LIT:20> , <NUM_LIT:30> , <NUM_LIT> , <NUM_LIT:50> ] , slice ( <NUM_LIT:1> , None , <NUM_LIT:2> ) ) == [ <NUM_LIT:20> , <NUM_LIT> ] <EOL> def test_hard_fuse_slice_cases ( ) : <EOL> term = ( getarray , ( getarray , '<STR_LIT:x>' , ( None , slice ( None , None ) ) ) , <EOL> ( slice ( None , None ) , <NUM_LIT:5> ) ) <EOL> assert rewrite_rules . rewrite ( term ) == ( getarray , '<STR_LIT:x>' , ( None , <NUM_LIT:5> ) ) </s>
