def _reduce(self, local_reduce_name, axes=None, dtype=None, out=None): if any(0 in localshape for localshape in self.localshapes()): raise NotImplementedError("Reduction not implemented for empty " "LocalArrays") if out is not None: _raise_nie() dtype = dtype or self.dtype out_dist = self.distribution.reduce(axes=axes) ddpr = out_dist.get_dim_data_per_rank() def _local_reduce(local_name, larr, out_comm, ddpr, dtype, axes): import distarray.localapi.localarray as la local_reducer = getattr(la, local_name) res = proxyize(la.local_reduction(out_comm, local_reducer, larr, # noqa ddpr, dtype, axes)) return res local_reduce_args = (local_reduce_name, self.key, out_dist.comm, ddpr, dtype, normalize_reduction_axes(axes, self.ndim)) out_key = self.context.apply(_local_reduce, local_reduce_args, targets=self.targets)[0] return DistArray.from_localarrays(key=out_key, distribution=out_dist, dtype=dtype)
def reduce(self, axes): """ Returns a new Distribution reduced along `axis`, i.e., the new distribution has one fewer dimension than `self`. """ # the `axis` argument can actually be a sequence of axes, so we rename it. axes = normalize_reduction_axes(axes, self.ndim) reduced_shape = remove_elements(axes, self.shape) reduced_dist = remove_elements(axes, self.dist) reduced_grid_shape = remove_elements(axes, self.grid_shape) # This block is required because np.min() works one axis at a time. reduced_ranks = self.rank_from_coords.copy() for axis in axes: reduced_ranks = np.min(reduced_ranks, axis=axis, keepdims=True) reduced_targets = [self.targets[r] for r in reduced_ranks.flat] return Distribution(context=self.context, shape=reduced_shape, dist=reduced_dist, grid_shape=reduced_grid_shape, targets=reduced_targets)
def _reduce(self, local_reduce_name, axes=None, dtype=None, out=None): if any(0 in localshape for localshape in self.localshapes()): raise NotImplementedError("Reduction not implemented for empty " "LocalArrays") if out is not None: _raise_nie() dtype = dtype or self.dtype out_dist = self.distribution.reduce(axes=axes) ddpr = out_dist.get_dim_data_per_rank() def _local_reduce(local_name, larr, out_comm, ddpr, dtype, axes): import distarray.localapi.localarray as la local_reducer = getattr(la, local_name) res = proxyize( la.local_reduction( out_comm, local_reducer, larr, # noqa ddpr, dtype, axes)) return res local_reduce_args = (local_reduce_name, self.key, out_dist.comm, ddpr, dtype, normalize_reduction_axes(axes, self.ndim)) out_key = self.context.apply(_local_reduce, local_reduce_args, targets=self.targets)[0] return DistArray.from_localarrays(key=out_key, distribution=out_dist, dtype=dtype)