<s> """<STR_LIT>""" <EOL> import sys <EOL> from six . moves import range , zip <EOL> import numpy as np <EOL> from numpy import ndarray <EOL> from itertools import product <EOL> def array_idx_iter ( shape ) : <EOL> """<STR_LIT>""" <EOL> for p in product ( * [ range ( s ) for s in shape ] ) : <EOL> yield p <EOL> def evenly_distrib_idxs ( num_divisions , arr_size ) : <EOL> """<STR_LIT>""" <EOL> base = arr_size // num_divisions <EOL> leftover = arr_size % num_divisions <EOL> sizes = np . ones ( num_divisions , dtype = "<STR_LIT:int>" ) * base <EOL> sizes [ : leftover ] += <NUM_LIT:1> <EOL> offsets = np . zeros ( num_divisions , dtype = "<STR_LIT:int>" ) <EOL> offsets [ <NUM_LIT:1> : ] = np . cumsum ( sizes ) [ : - <NUM_LIT:1> ] <EOL> return sizes , offsets <EOL> def to_slice ( idxs ) : <EOL> """<STR_LIT>""" <EOL> if len ( idxs ) == <NUM_LIT:1> : <EOL> return slice ( idxs [ <NUM_LIT:0> ] , idxs [ <NUM_LIT:0> ] + <NUM_LIT:1> ) <EOL> elif len ( idxs ) == <NUM_LIT:0> : <EOL> return idxs <EOL> stride = idxs [ <NUM_LIT:1> ] - idxs [ <NUM_LIT:0> ] <EOL> if stride <= <NUM_LIT:0> : <EOL> return idxs <EOL> if any ( idxs [ <NUM_LIT:1> : ] - idxs [ : - <NUM_LIT:1> ] != stride ) : <EOL> return idxs <EOL> return slice ( idxs [ <NUM_LIT:0> ] , idxs [ - <NUM_LIT:1> ] + stride , stride ) </s>
