diff --git a/distarray/dist/tests/test_distarray.py b/distarray/dist/tests/test_distarray.py index 2ffd7c9a..8c666a1c 100644 --- a/distarray/dist/tests/test_distarray.py +++ b/distarray/dist/tests/test_distarray.py @@ -595,6 +595,12 @@ def test_sum_4D_cyclic(self): assert_allclose(darr_sum.tondarray(), arr_sum) assert_allclose(darr.sum().tondarray(), arr.sum()) + def test_gh_435_regression_with_var(self): + dist = Distribution.from_shape(self.context, shape=(14,), dist=('b'), + targets=range(4)) + darr = self.context.ones(dist) + darr.var() + class TestFromLocalArrays(ContextTestCase): diff --git a/distarray/local/localarray.py b/distarray/local/localarray.py index b8dc3ea2..10eb9875 100644 --- a/distarray/local/localarray.py +++ b/distarray/local/localarray.py @@ -950,23 +950,27 @@ def _basic_reducer(reduce_comm, op, func, args, kwargs, out): reduce_comm.Reduce(local_reduce, out_ndarray, op=op, root=0) return out + def min_reducer(reduce_comm, larr, out, axes, dtype): """ Core reduction function for min.""" return _basic_reducer(reduce_comm, MPI.MIN, - larr.ndarray.min, - (), {'axis':axes}, out) + larr.ndarray.min, + (), {'axis': axes}, out) + def max_reducer(reduce_comm, larr, out, axes, dtype): """ Core reduction function for max.""" return _basic_reducer(reduce_comm, MPI.MAX, - larr.ndarray.max, - (), {'axis':axes}, out) + larr.ndarray.max, + (), {'axis': axes}, out) + def sum_reducer(reduce_comm, larr, out, axes, dtype): """ Core reduction function for sum.""" return _basic_reducer(reduce_comm, MPI.SUM, - larr.ndarray.sum, - (), {'axis':axes, 'dtype':dtype}, out) + larr.ndarray.sum, + (), {'axis': axes, 'dtype': dtype}, out) + def mean_reducer(reduce_comm, larr, out, axes, dtype): """ Core reduction function for mean.""" @@ -978,28 +982,27 @@ def mean_reducer(reduce_comm, larr, out, axes, dtype): def var_reducer(reduce_comm, larr, out, axes, dtype): """ Core reduction function for var.""" + temp = empty_like(larr, dtype=float) + + # We hold the intermediate means in `mean`. mean = empty_like(out, dtype=float) if out is not None else None mean = mean_reducer(reduce_comm, larr, mean, axes, dtype=float) - temp = empty_like(larr, dtype=float) + # Have to broadcast mean.ndarray to all ranks in this reduce_comm. + mean = reduce_comm.bcast(mean, root=0) - # Make mean.ndarray's shape broadcastable. - if mean is not None: - mean_shape = tuple(1 if axis in axes else s - for (axis, s) in enumerate(larr.ndarray.shape)) - mean.ndarray.shape = mean_shape - # Copy mean.ndarray into temp.ndarray - temp.ndarray[...] = mean.ndarray + mean_shape = [] + for (ax, s) in enumerate(larr.ndarray.shape): + mean_shape.append(1 if ax in axes else s) + mean_shape = tuple(mean_shape) - # have to broadcast mean.ndarray to all ranks in this reduce_comm. - reduce_comm.Bcast(temp.ndarray, root=0) + mean.ndarray.shape = mean_shape # Do the variance calculation. - temp.ndarray[...] = (larr.ndarray - temp.ndarray) ** 2 + temp.ndarray[...] = (larr.ndarray - mean.ndarray) ** 2 # Get the mean reduction of temp's data. mean_reducer(reduce_comm, temp, out, axes, dtype) - return out