def aggregate_shape(leaf, expr, chunk, chunk_expr):
""" The shape of the intermediate aggregate
>>> leaf = symbol('leaf', '10 * 10 * int')
>>> expr = leaf.sum(axis=0)
>>> chunk = symbol('chunk', '3 * 3 * int') # 3 does not divide 10
>>> chunk_expr = chunk.sum(axis=0, keepdims=True)
>>> aggregate_shape(leaf, expr, chunk, chunk_expr)
(4, 10)
"""
if datashape.var in tconcat(map(shape, [leaf, expr, chunk, chunk_expr])):
return (datashape.var, ) * leaf.ndim
numblocks = [int(floor(l / c)) for l, c in zip(leaf.shape, chunk.shape)]
last_chunk_shape = [l % c for l, c in zip(leaf.shape, chunk.shape)]
if builtins.sum(last_chunk_shape) != 0:
old = last_chunk_shape
last_chunk = symbol(chunk._name,
DataShape(*(last_chunk_shape + [chunk.dshape.measure])))
last_chunk_expr = chunk_expr._subs({chunk: last_chunk})
last_chunk_shape = shape(last_chunk_expr)
# Keep zeros if they were there before
last_chunk_shape = tuple(a if b != 0 else 0 for a, b in
zip(last_chunk_shape, old))
return tuple(int(floor(l / c)) * ce + lce
for l, c, ce, lce
in zip(shape(leaf), shape(chunk), shape(chunk_expr), last_chunk_shape))
def aggregate_shape(leaf, expr, chunk, chunk_expr):
""" The shape of the intermediate aggregate
>>> leaf = symbol('leaf', '10 * 10 * int')
>>> expr = leaf.sum(axis=0)
>>> chunk = symbol('chunk', '3 * 3 * int') # 3 does not divide 10
>>> chunk_expr = chunk.sum(axis=0, keepdims=True)
>>> aggregate_shape(leaf, expr, chunk, chunk_expr)
(4, 10)
"""
if datashape.var in tconcat(map(shape, [leaf, expr, chunk, chunk_expr])):
return (datashape.var, ) * leaf.ndim
numblocks = [int(floor(l / c)) for l, c in zip(leaf.shape, chunk.shape)]
last_chunk_shape = [l % c for l, c in zip(leaf.shape, chunk.shape)]
if builtins.sum(last_chunk_shape) != 0:
old = last_chunk_shape
last_chunk = symbol(
chunk._name,
DataShape(*(last_chunk_shape + [chunk.dshape.measure])))
last_chunk_expr = chunk_expr._subs({chunk: last_chunk})
last_chunk_shape = shape(last_chunk_expr)
# Keep zeros if they were there before
last_chunk_shape = tuple(a if b != 0 else 0
for a, b in zip(last_chunk_shape, old))
return tuple(
int(floor(l / c)) * ce + lce for l, c, ce, lce in zip(
shape(leaf), shape(chunk), shape(chunk_expr), last_chunk_shape))
def _leaves(self):
return list(unique(tconcat(i._leaves() for i in self.children)))
def fields(self):
return list(tconcat(child.fields for child in self.children))
def _leaves(self):
return list(unique(tconcat(i._leaves() for i in self.children)))
def fields(self):
return list(tconcat(child.fields for child in self.children))
def fields(self):
return list(tconcat(arg.fields for arg in self.args))