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