def test_global_flat_indices(self): dist0 = Distribution(self.context, (40, ), ('b', ), (2, ), targets=[1, 3]) self.assertSequenceEqual([[(0, 20)], [(20, 40)]], [ global_flat_indices(ddpr) for ddpr in dist0.get_dim_data_per_rank() ]) dist1 = Distribution(self.context, (5, 8), ('b', 'n'), (2, ), targets=[0, 2]) self.assertSequenceEqual([[(0, 24)], [(24, 40)]], [ global_flat_indices(ddpr) for ddpr in dist1.get_dim_data_per_rank() ]) dist2 = Distribution(self.context, (5, 8), ('b', 'b'), (2, 2), targets=[0, 1, 2, 3]) self.assertSequenceEqual( [[(0, 4), (8, 12), (16, 20)], [(4, 8), (12, 16), (20, 24)], [(24, 28), (32, 36)], [(28, 32), (36, 40)]], [ global_flat_indices(ddpr) for ddpr in dist2.get_dim_data_per_rank() ])
def test_global_flat_indices(self): dist0 = Distribution(self.context, (40,), ("b",), (2,), targets=[1, 3]) self.assertSequenceEqual( [[(0, 20)], [(20, 40)]], [global_flat_indices(ddpr) for ddpr in dist0.get_dim_data_per_rank()] ) dist1 = Distribution(self.context, (5, 8), ("b", "n"), (2,), targets=[0, 2]) self.assertSequenceEqual( [[(0, 24)], [(24, 40)]], [global_flat_indices(ddpr) for ddpr in dist1.get_dim_data_per_rank()] ) dist2 = Distribution(self.context, (5, 8), ("b", "b"), (2, 2), targets=[0, 1, 2, 3]) self.assertSequenceEqual( [[(0, 4), (8, 12), (16, 20)], [(4, 8), (12, 16), (20, 24)], [(24, 28), (32, 36)], [(28, 32), (36, 40)]], [global_flat_indices(ddpr) for ddpr in dist2.get_dim_data_per_rank()], )
def test_create_target_subset(self): shape = (100, 100) subtargets = self.context.targets[::2] distribution = Distribution(self.context, shape=shape, targets=subtargets) darr = self.context.ones(distribution) lss = darr.localshapes() self.assertEqual(len(lss), len(subtargets)) ddpr = distribution.get_dim_data_per_rank() self.assertEqual(len(ddpr), len(subtargets))
def fromfunction(self, function, shape, **kwargs): """Create a DistArray from a function over global indices. Unlike numpy's `fromfunction`, the result of distarray's `fromfunction` is restricted to the same Distribution as the index array generated from `shape`. See numpy.fromfunction for more details. """ self.push_function(function.__name__, function, targets=self.targets) def _local_fromfunction(func_name, comm, ddpr, kwargs): from distarray.localapi import fromfunction from distarray.localapi.maps import Distribution from importlib import import_module main = import_module('__main__') if len(ddpr): dim_data = ddpr[comm.Get_rank()] else: dim_data = () func = getattr(main, func_name) dist = Distribution(comm, dim_data=dim_data) local_arr = fromfunction(func, dist, **kwargs) return proxyize(local_arr) dist = kwargs.get('dist', None) grid_shape = kwargs.get('grid_shape', None) distribution = Distribution(context=self, shape=shape, dist=dist, grid_shape=grid_shape) ddpr = distribution.get_dim_data_per_rank() da_name = self.apply( _local_fromfunction, (function.__name__, distribution.comm, ddpr, kwargs), targets=distribution.targets) return DistArray.from_localarrays(da_name[0], distribution=distribution)
def fromfunction(self, function, shape, **kwargs): """Create a DistArray from a function over global indices. Unlike numpy's `fromfunction`, the result of distarray's `fromfunction` is restricted to the same Distribution as the index array generated from `shape`. See numpy.fromfunction for more details. """ self.push_function(function.__name__, function, targets=self.targets) def _local_fromfunction(func_name, comm, ddpr, kwargs): from distarray.localapi import fromfunction from distarray.localapi.maps import Distribution from importlib import import_module main = import_module('__main__') if len(ddpr): dim_data = ddpr[comm.Get_rank()] else: dim_data = () func = getattr(main, func_name) dist = Distribution(comm, dim_data=dim_data) local_arr = fromfunction(func, dist, **kwargs) return proxyize(local_arr) dist = kwargs.get('dist', None) grid_shape = kwargs.get('grid_shape', None) distribution = Distribution(context=self, shape=shape, dist=dist, grid_shape=grid_shape) ddpr = distribution.get_dim_data_per_rank() da_name = self.apply(_local_fromfunction, (function.__name__, distribution.comm, ddpr, kwargs), targets=distribution.targets) return DistArray.from_localarrays(da_name[0], distribution=distribution)