Fix DistArray.var with different sized localarrays (closes issue 435)

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

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