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