Merge pull request #396 from cowlicks/recurs-bug

Fix bug wth all 'n's dist tuple.

Author: kwmsmith

distarray/dist/maps.py

@@ -36,7 +36,6 @@
                                       normalize_grid_shape,
                                       make_grid_shape,
                                       positivify,
-                                      validate_grid_shape,
                                       _start_stop_block,
                                       normalize_dim_dict,
                                       normalize_reduction_axes)
@@ -412,8 +411,8 @@ def from_dim_data_per_rank(cls, context, dim_data_per_rank, targets=None):
         self.ndim = len(dd0)
         self.dist = tuple(dd['dist_type'] for dd in dd0)
         self.grid_shape = tuple(dd['proc_grid_size'] for dd in dd0)
-
-        validate_grid_shape(self.grid_shape, self.dist, len(self.targets))
+        self.grid_shape = normalize_grid_shape(self.grid_shape, self.ndim,
+                                               self.dist, len(self.targets))
 
         coords = [tuple(d['proc_grid_rank'] for d in dd) for dd in
                   dim_data_per_rank]
@@ -439,7 +438,18 @@ def from_dim_data_per_rank(cls, context, dim_data_per_rank, targets=None):
         return self
 
     @classmethod
-    def from_shape(cls, context, shape, dist=None, grid_shape=None, targets=None):
+    def from_shape(cls, context, shape, dist=None, grid_shape=None,
+                   targets=None):
+
+        # special case when dist is all 'n's.
+        if (dist is not None) and all(d == 'n' for d in dist):
+            if (targets is not None) and (len(targets) != 1):
+                raise ValueError('target dist conflict')
+            elif targets is None:
+                targets = [context.targets[0]]
+            else:
+                # then targets is set correctly
+                pass
 
         self = cls.__new__(cls)
         self.context = context
@@ -448,17 +458,18 @@ def from_shape(cls, context, shape, dist=None, grid_shape=None, targets=None):
         self.shape = shape
         self.ndim = len(shape)
 
+        # dist
         if dist is None:
             dist = {0: 'b'}
         self.dist = normalize_dist(dist, self.ndim)
 
-        if grid_shape is None:  # Make a new grid_shape if not provided.
-            self.grid_shape = make_grid_shape(self.shape, self.dist,
-                                              len(self.targets))
-        else:  # Otherwise normalize the one passed in.
-            self.grid_shape = normalize_grid_shape(grid_shape, self.ndim)
-        # In either case, validate.
-        validate_grid_shape(self.grid_shape, self.dist, len(self.targets))
+        # grid_shape
+        if grid_shape is None:
+            grid_shape = make_grid_shape(self.shape, self.dist,
+                                         len(self.targets))
+
+        self.grid_shape = normalize_grid_shape(grid_shape, self.ndim,
+                                               self.dist, len(self.targets))
 
         # TODO: FIXME: assert that self.rank_from_coords is valid and conforms
         # to how MPI does it.
@@ -568,7 +579,8 @@ def __init__(self, context, global_dim_data, targets=None):
         self.dist = tuple(m.dist for m in self.maps)
         self.grid_shape = tuple(m.grid_size for m in self.maps)
 
-        validate_grid_shape(self.grid_shape, self.dist, len(context.targets))
+        self.grid_shape = normalize_grid_shape(self.grid_shape, self.ndim,
+                                               self.dist, len(context.targets))
 
         nelts = reduce(operator.mul, self.grid_shape, 1)
         self.rank_from_coords = np.arange(nelts).reshape(self.grid_shape)

distarray/dist/tests/test_maps.py

@@ -71,7 +71,7 @@ def test_is_compatible(self):
 
         self.assertTrue(cm0.is_compatible(cm1))
         self.assertTrue(cm1.is_compatible(cm0))
-        
+
         nr -= 1; nc -= 1; nd -= 1
 
         cm2 = client_map.Distribution.from_shape(
@@ -112,3 +112,15 @@ def test_reduce_0D(self):
         self.assertSequenceEqual(new_dist.shape, ())
         self.assertEqual(new_dist.grid_shape, ())
         self.assertEqual(set(new_dist.targets), set(dist.targets[:1]))
+
+
+class TestDistributionCreation(unittest.TestCase):
+    def test_all_n_dist(self):
+        context = Context()
+        distribution = client_map.Distribution.from_shape(context,
+                                                          shape=(3, 3),
+                                                          dist=('n', 'n'))
+        context.ones(distribution)
+
+if __name__ == '__main__':
+    unittest.main(verbosity=2)

distarray/local/maps.py

@@ -27,9 +27,9 @@
 from distarray.externals.six.moves import range, zip
 
 from distarray.local import construct
-from distarray.metadata_utils import (validate_grid_shape, make_grid_shape,
-                                      normalize_grid_shape, normalize_dist,
-                                      distribute_indices, positivify)
+from distarray.metadata_utils import (make_grid_shape, normalize_grid_shape,
+                                      normalize_dist, distribute_indices,
+                                      positivify)
 
 
 class Distribution(object):
@@ -56,10 +56,8 @@ def from_shape(cls, shape, dist=None, grid_shape=None, comm=None):
 
         if grid_shape is None:  # Make a new grid_shape if not provided.
             grid_shape = make_grid_shape(shape, dist_tuple, comm_size)
-        else:  # Otherwise normalize the one passed in.
-            grid_shape = normalize_grid_shape(grid_shape, ndim)
-        # In either case, validate.
-        validate_grid_shape(grid_shape, dist_tuple, comm_size)
+        grid_shape = normalize_grid_shape(grid_shape, ndim,
+                                          dist_tuple, comm_size)
 
         comm = construct.init_comm(base_comm, grid_shape)
         grid_coords = comm.Get_coords(comm.Get_rank())
@@ -284,7 +282,7 @@ class BlockCyclicMap(MapBase):
     """
 
     dist = 'c'
-    
+
     def __init__(self, global_size, grid_size, grid_rank, start, block_size):
         if start % block_size:
             msg = "Value of start (%r) does not evenly divide block_size (%r)."

distarray/local/tests/paralleltest_maps.py

@@ -217,7 +217,7 @@ def test_grid_shape(self):
         dist = Distribution.from_shape((20, 20), dist='b', comm=self.comm)
         self.assertEqual(dist.grid_shape, (12, 1))
         dist = Distribution.from_shape((2*10, 6*10), dist=('b', 'b'),
-                                                           comm=self.comm)
+                                       comm=self.comm)
         self.assertEqual(dist.grid_shape, (2, 6))
         dist = Distribution.from_shape((6*10, 2*10), dist='bb', comm=self.comm)
         self.assertEqual(dist.grid_shape, (6, 2))

distarray/metadata_utils.py

@@ -24,15 +24,19 @@ class GridShapeError(Exception):
     pass
 
 
-def normalize_grid_shape(grid_shape, ndims):
-    """Adds 1s to grid_shape so it has `ndims` dimensions."""
-    return tuple(grid_shape) + (1,) * (ndims - len(grid_shape))
+def normalize_grid_shape(grid_shape, ndims, dist, comm_size):
+    """Adds 1s to grid_shape so it has `ndims` dimensions.  Validates
+    `grid_shape` tuple against the `dist` tuple and `comm_size`.
+    """
+    grid_shape = tuple(grid_shape) + (1,) * (ndims - len(grid_shape))
 
+    # short circuit for special case
+    if all(x == 'n' for x in dist):
+        if not all(x == 1 for x in grid_shape):
+            raise ValueError("grid shape should be all `1`'s not %s." %
+                             grid_shape)
+        return grid_shape
 
-def validate_grid_shape(grid_shape, dist, comm_size):
-    """ Validates `grid_shape` tuple against the `dist` tuple and
-    `comm_size`.
-    """
     if len(grid_shape) != len(dist):
         msg = "grid_shape's length (%d) not equal to dist's length (%d)"
         raise InvalidGridShapeError(msg % (len(grid_shape), len(dist)))
@@ -74,7 +78,9 @@ def make_grid_shape(shape, dist, comm_size):
     distdims = tuple(i for (i, v) in enumerate(dist) if v != 'n')
     ndistdim = len(distdims)
 
-    if ndistdim == 1:
+    if ndistdim == 0:
+        dist_grid_shape = ()
+    elif ndistdim == 1:
         # Trivial case: all processes used for the one distributed dimension.
         dist_grid_shape = (comm_size,)
 
@@ -223,4 +229,3 @@ def normalize_reduction_axes(axes, ndim):
     else:
         axes = tuple(positivify(a, ndim) for a in axes)
     return axes
-