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, 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)
