try: from blaze.compute.numba import get_numba_ufunc, broadcast_collect, Broadcastable def compute_broadcast(expr, *data, **kwargs): expr_inds = tuple(range(ndim(expr)))[::-1] func = get_numba_ufunc(expr) return atop(func, expr_inds, *concat((dat, tuple(range(ndim(dat))[::-1])) for dat in data)) def optimize_array(expr, *data): return broadcast_collect(expr, Broadcastable=Broadcastable, WantToBroadcast=Broadcastable) for i in range(5): compute_up.register(Broadcast, *([(Array, Number)] * i))(compute_broadcast) optimize.register(Expr, *([(Array, Number)] * i))(optimize_array) except ImportError: pass for i in range(5): compute_up.register(ElemWise, *([Array] * i))(elemwise_array) @dispatch(Reduction, Array) def compute_up(expr, data, **kwargs): leaf = expr._leaves()[0] chunk = symbol("chunk", DataShape(*(tuple(map(first, data.chunks)) + (leaf.dshape.measure,)))) (chunk, chunk_expr), (agg, agg_expr) = split(expr._child, expr, chunk=chunk)
def compute_broadcast(expr, *data, **kwargs): expr_inds = tuple(range(ndim(expr)))[::-1] func = get_numba_ufunc(expr) return atop( func, next(names), expr_inds, *concat((dat, tuple(range(ndim(dat))[::-1])) for dat in data)) def optimize_array(expr, *data): return broadcast_collect(expr, Broadcastable=Broadcastable, WantToBroadcast=Broadcastable) for i in range(5): compute_up.register(Broadcast, *([(Array, Number)] * i))(compute_broadcast) optimize.register(Expr, *([(Array, Number)] * i))(optimize_array) except ImportError: pass for i in range(5): compute_up.register(ElemWise, *([Array] * i))(elemwise_array) @dispatch(Reduction, Array) def compute_up(expr, data, **kwargs): leaf = expr._leaves()[0] chunk = symbol( 'chunk', DataShape(*(tuple(map(first, data.chunks)) + (leaf.dshape.measure, )))) (chunk, chunk_expr), (agg, agg_expr) = split(expr._child,