def reduce_axis(self, op_name, axis, keepdims=False):
result_blocks = np.empty_like(self.blocks, dtype=Block)
for grid_entry in self.grid.get_entry_iterator():
result_blocks[grid_entry] = self.blocks[grid_entry].reduce_axis(op_name,
axis,
keepdims=keepdims)
result_shape = []
result_block_shape = []
for curr_axis in range(len(self.shape)):
axis_size, axis_block_size = self.shape[curr_axis], self.block_shape[curr_axis]
if curr_axis == axis:
if keepdims:
axis_size, axis_block_size = 1, 1
else:
continue
result_shape.append(axis_size)
result_block_shape.append(axis_block_size)
result_shape = tuple(result_shape)
result_block_shape = tuple(result_block_shape)
result_dtype = array_utils.get_reduce_output_type(op_name, self.dtype)
result_grid = ArrayGrid(shape=result_shape,
block_shape=result_block_shape,
dtype=result_dtype.__name__)
result = BlockArray(result_grid, self.system)
op_func = np.__getattribute__(op_name)
reduced_blocks = op_func(result_blocks, axis=axis, keepdims=keepdims)
if result.shape == ():
result.blocks[()] = reduced_blocks
else:
result.blocks = reduced_blocks
return result
def bop_reduce(self, op_name, other):
other: Block = other
dtype = array_utils.get_reduce_output_type(op_name, self.dtype)
block = Block(
grid_entry=self.grid_entry,
grid_shape=self.grid_shape,
rect=list(self.rect),
shape=self.shape,
dtype=dtype,
# This is false because we invoke the transpose before
# applying the reduction operation, so the resulting
# remote object will be transposed.
transposed=False,
system=self._system)
block.oid = self._system.bop_reduce(op=op_name,
a1=self.oid,
a2=other.oid,
a1_T=self.transposed,
a2_T=other.transposed,
syskwargs={
"grid_entry": block.grid_entry,
"grid_shape": block.grid_shape
})
return block
def reduce_axis(self, op_name, axis, keepdims):
# TODO (hme): Add options version of this too, but need to fix block imps
# so we're not branching between options / no options calls in every method.
# We lose an axis, so make sure to account for dropped axis in result.
result_grid_entry = []
result_grid_shape = []
result_rect = []
result_shape = []
for curr_axis in range(len(self.shape)):
if curr_axis == axis or axis is None:
if keepdims:
result_grid_entry.append(0)
result_grid_shape.append(1)
result_rect.append((0, 1))
result_shape.append(1)
continue
result_grid_entry.append(self.grid_entry[curr_axis])
result_grid_shape.append(self.grid_shape[curr_axis])
result_rect.append(self.rect[curr_axis])
result_shape.append(self.shape[curr_axis])
dtype = array_utils.get_reduce_output_type(op_name, self.dtype)
block = Block(
grid_entry=result_grid_entry,
grid_shape=result_grid_shape,
rect=result_rect,
shape=result_shape,
dtype=dtype,
# This is false because we invoke the transpose before
# applying the reduction operation, so the resulting
# remote object will be transposed.
transposed=False,
system=self._system)
block.oid = self._system.reduce_axis(op_name=op_name,
arr=self.oid,
axis=axis,
keepdims=keepdims,
transposed=self.transposed,
syskwargs={
"grid_entry":
block.grid_entry,
"grid_shape": block.grid_shape
})
return block
def reduce_axis(self, op_name, axis, keepdims=False):
if not (axis is None or isinstance(axis, (int, np.int32, np.int64))):
raise NotImplementedError("Only integer axis is currently supported.")
result_blocks = np.empty_like(self.blocks, dtype=Block)
for grid_entry in self.grid.get_entry_iterator():
result_blocks[grid_entry] = self.blocks[grid_entry].reduce_axis(op_name,
axis,
keepdims=keepdims)
result_shape = []
result_block_shape = []
for curr_axis in range(len(self.shape)):
axis_size, axis_block_size = self.shape[curr_axis], self.block_shape[curr_axis]
if curr_axis == axis or axis is None:
if keepdims:
axis_size, axis_block_size = 1, 1
else:
continue
result_shape.append(axis_size)
result_block_shape.append(axis_block_size)
result_shape = tuple(result_shape)
result_block_shape = tuple(result_block_shape)
result_dtype = array_utils.get_reduce_output_type(op_name, self.dtype)
result_grid = ArrayGrid(shape=result_shape,
block_shape=result_block_shape,
dtype=result_dtype.__name__)
result = BlockArray(result_grid, self.system)
if op_name in settings.np_pairwise_reduction_map:
# Do a pairwise reduction with the pairwise reduction op.
pairwise_op_name = settings.np_pairwise_reduction_map.get(op_name, op_name)
if axis is None:
reduced_block: Block = None
for grid_entry in self.grid.get_entry_iterator():
if reduced_block is None:
reduced_block = result_blocks[grid_entry]
continue
next_block = result_blocks[grid_entry]
reduced_block = reduced_block.bop(pairwise_op_name, next_block, {})
if result.shape == ():
result.blocks[()] = reduced_block
else:
result.blocks[:] = reduced_block
else:
for result_grid_entry in result_grid.get_entry_iterator():
reduced_block: Block = None
for sum_dim in range(self.grid.grid_shape[axis]):
grid_entry = list(result_grid_entry)
if keepdims:
grid_entry[axis] = sum_dim
else:
grid_entry = grid_entry[:axis] + [sum_dim] + grid_entry[axis:]
grid_entry = tuple(grid_entry)
next_block: Block = result_blocks[grid_entry]
if reduced_block is None:
reduced_block = next_block
else:
reduced_block = reduced_block.bop(pairwise_op_name, next_block, {})
result.blocks[result_grid_entry] = reduced_block
else:
op_func = np.__getattribute__(op_name)
if result.shape == ():
result.blocks[()] = op_func(result_blocks, axis=axis, keepdims=keepdims)
else:
result.blocks = op_func(result_blocks, axis=axis, keepdims=keepdims)
return result
def reduce_axis(self, op_name, axis, keepdims=False):
if not (axis is None or isinstance(axis, (int, np.int32, np.int64))):
raise NotImplementedError(
"Only integer axis is currently supported.")
block_reduced_oids = np.empty_like(self.blocks, dtype=tuple)
for grid_entry in self.grid.get_entry_iterator():
block = self.blocks[grid_entry]
block_oid = self.cm.reduce_axis(
op_name=op_name,
arr=block.oid,
axis=axis,
keepdims=keepdims,
transposed=block.transposed,
syskwargs={
"grid_entry": block.grid_entry,
"grid_shape": block.grid_shape,
},
)
block_reduced_oids[grid_entry] = (
block_oid,
block.grid_entry,
block.grid_shape,
False,
)
result_shape = []
result_block_shape = []
for curr_axis in range(len(self.shape)):
axis_size, axis_block_size = (
self.shape[curr_axis],
self.block_shape[curr_axis],
)
if curr_axis == axis or axis is None:
if keepdims:
axis_size, axis_block_size = 1, 1
else:
continue
result_shape.append(axis_size)
result_block_shape.append(axis_block_size)
result_shape = tuple(result_shape)
result_block_shape = tuple(result_block_shape)
result_dtype = array_utils.get_reduce_output_type(op_name, self.dtype)
result_grid = ArrayGrid(
shape=result_shape,
block_shape=result_block_shape,
dtype=result_dtype.__name__,
)
result = BlockArray(result_grid, self.cm)
if axis is None:
if result.shape == ():
result_block: Block = result.blocks[()]
else:
result_block: Block = result.blocks[:].item()
result_block.oid = self._tree_reduce(
op_name,
block_reduced_oids.flatten().tolist(),
result_block.grid_entry,
result_block.grid_shape,
)
else:
for result_grid_entry in result_grid.get_entry_iterator():
block_reduced_oids_axis = []
for sum_dim in range(self.grid.grid_shape[axis]):
grid_entry = list(result_grid_entry)
if keepdims:
grid_entry[axis] = sum_dim
else:
grid_entry = grid_entry[:axis] + [sum_dim
] + grid_entry[axis:]
grid_entry = tuple(grid_entry)
block_reduced_oids_axis.append(
block_reduced_oids[grid_entry])
result_block: Block = result.blocks[result_grid_entry]
result_block.oid = self._tree_reduce(
op_name,
block_reduced_oids_axis,
result_block.grid_entry,
result_block.grid_shape,
)
return result
