New da.tile implementation

Someone privately sent me this da.tile implementation.

def tile(A, reps):
    try:
        tup = tuple(reps)
    except TypeError:
        tup = (reps,)
    d = len(tup)
    if (d < A.ndim):
        tup = (1,)*(A.ndim-d) + tup
    nnd = len(tup) - A.ndim # Number of new dimensions
    final_shape = tup[:nnd]+tuple(i*di for i,di in zip(tup[nnd:],A.shape))
    if 0 in final_shape:
        XXXXX # how to make empty array?
    tup = tup[nnd:]
    # Don't tile along dimensions that can be broadcasted
    tup = tuple(1 if l == 1 else r for l, r in zip(A.shape, tup))
    if all(i==1 for i in tup):
        return da.broadcast_to(A, final_shape)
    name = 'tile-'+dask.base.tokenize('tile', A.name, tup)
    dsk = dict(zip(itertools.product([name], *tuple(range(r*len(c)) for r,c in zip(tup, A.chunks))),
                   itertools.product([A.name], *tuple(r*range(len(c)) for r,c in zip(tup, A.chunks)))))
    ary =  da.Array(sharedict.merge(A.dask, dsk), name, tuple(r*c for r,c in zip(tup, A.chunks)), A.dtype)
    return da.broadcast_to(ary, final_shape)

I haven’t yet gone through it. If anyone feels like an easy contribution this could be checked and tests could be added

[mrocklin]

mrocklin@carbon:~/workspace/dask$ git grep "def tile"
dask/array/creation.py:def tile(A, reps):

[jakirkham]

FWIW one could probably solve this issue today by implementing da.tile using da.block.

In [1]: import dask.array as da                                                 

In [2]: a = da.asarray([0, 1, 2])                                               

In [3]: da.block(3 * [4 * [a]]).compute()                                       
Out[3]: 
array([[0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2],
       [0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2],
       [0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2]])