Adds workaround to remove empty local arrays.

Slices the distribution object to remove all empty localarrays before
returning.  Only done if all dimensions use 'b' or 'n' in the dist
tuple.

Author: Kurt Smith

distarray/dist/maps.py

@@ -612,7 +612,17 @@ def __new__(cls, context, shape, dist=None, grid_shape=None, targets=None):
 
         # list of `ClientMap` objects, one per dimension.
         maps = [map_from_sizes(*args) for args in zip(shape, dist, grid_shape)]
-        return cls.from_maps(context=context, maps=maps, targets=targets)
+
+        self = cls.from_maps(context=context, maps=maps, targets=targets)
+
+        # TODO: FIXME: this is a workaround.  The reason we slice here is to
+        # return a distribution with no empty local shapes.  The `from_maps()`
+        # classmethod should be fixed to ensure no empty local arrays are
+        # created in the first place.  That will remove the need to slice the
+        # distribution to remove empty localshapes.
+        if all(d in ('n', 'b') for d in self.dist):
+            self = self.slice((slice(None),)*self.ndim)
+        return self
 
     @classmethod
     def from_global_dim_data(cls, context, global_dim_data, targets=None):

distarray/dist/tests/test_maps.py

@@ -300,5 +300,23 @@ def test_all_n_dist(self):
         self.context.ones(distribution)
 
 
+class TestNoEmptyLocals(ContextTestCase):
+
+    def test_no_empty_local_arrays_4_targets(self):
+        for n in range(1, 20):
+            dist = Distribution(self.context, shape=(n,),
+                                dist=('b',),
+                                targets=self.context.targets[:4])
+            for ls in dist.localshapes():
+                self.assertNotIn(0, ls)
+
+    def test_no_empty_local_arrays_3_targets(self):
+        for n in range(1, 20):
+            dist = Distribution(self.context, shape=(n,),
+                                dist=('b',),
+                                targets=self.context.targets[:3])
+            for ls in dist.localshapes():
+                self.assertNotIn(0, ls)
+
 if __name__ == '__main__':
     unittest.main(verbosity=2)