def test_block_redist_2D_one_to_many(self):
source_dist = Distribution(self.context, (9, 9), ('b', 'b'), (1, 1),
targets=[2])
dest_dist = Distribution(self.context, (9, 9), ('b', 'b'), (2, 2),
targets=range(4))
plan = source_dist.get_redist_plan(dest_dist)
expected = [
{
'source_rank': 2,
'dest_rank': 0,
'indices': [(0, 5, 1), (0, 5, 1)]
},
{
'source_rank': 2,
'dest_rank': 1,
'indices': [(0, 5, 1), (5, 9, 1)]
},
{
'source_rank': 2,
'dest_rank': 2,
'indices': [(5, 9, 1), (0, 5, 1)]
},
{
'source_rank': 2,
'dest_rank': 3,
'indices': [(5, 9, 1), (5, 9, 1)]
},
]
for p, e in zip(plan, expected):
self.assertEqual(p, e)
def test_not_compatible(self):
dist_b1 = Distribution(self.context, (10, ), ('b', ), (1, ),
targets=[0])
dist_b2 = Distribution(self.context, (9, ), ('b', ), (1, ),
targets=[0])
self.assertFalse(dist_b1.is_compatible(dist_b2))
self.assertFalse(dist_b2.is_compatible(dist_b1))
dist_b3 = Distribution(self.context, (10, ), ('b', ), (2, ),
targets=[0, 1])
self.assertFalse(dist_b1.is_compatible(dist_b3))
self.assertFalse(dist_b3.is_compatible(dist_b1))
dist_b4 = Distribution(self.context, (10, ), ('c', ), (2, ),
targets=[0, 1])
self.assertFalse(dist_b4.is_compatible(dist_b3))
self.assertFalse(dist_b3.is_compatible(dist_b4))
gdd_unstructured = ({
'dist_type': 'u',
'indices': [range(10)],
}, )
dist_u = Distribution.from_global_dim_data(self.context,
gdd_unstructured)
self.assertFalse(dist_u.is_compatible(dist_b1))
self.assertFalse(dist_b1.is_compatible(dist_u))
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_is_compatible_degenerate(self):
dist_bc = Distribution(self.context, (10, 10), ('b', 'c'), (1, 1),
targets=[0])
dist_cb = Distribution(self.context, (10, 10), ('c', 'b'), (1, 1),
targets=[0])
self.assertTrue(dist_bc.is_compatible(dist_cb))
self.assertTrue(dist_cb.is_compatible(dist_bc))
def test_is_compatible_nodist(self):
# See GH issue #461.
dist_bcn = Distribution(self.context, (10, 10, 10), ('b', 'c', 'n'),
(1, 1, 1),
targets=[0])
dist_nnn = Distribution(self.context, (10, 10, 10), ('n', 'n', 'n'),
(1, 1, 1),
targets=[0])
self.assertTrue(dist_bcn.is_compatible(dist_nnn))
self.assertTrue(dist_nnn.is_compatible(dist_bcn))
def test_3D(self):
a = self.context.empty((3, 3, 3))
db = Distribution(self.context, (3, 3, 3),
dist=('n', 'b', 'n'))
b = self.context.empty(db)
dc = Distribution(self.context, (3, 3, 3),
dist=('n', 'n', 'b'))
c = self.context.empty(dc)
self.assertEqual(a.grid_shape, (3, 1, 1))
self.assertEqual(b.grid_shape, (1, 3, 1))
self.assertEqual(c.grid_shape, (1, 1, 3))
def test_redist_incompatible_sizes(self):
source_da = self.context.empty((10, ), dtype=numpy.int32)
with self.assertRaises(ValueError):
source_da.distribute_as(Distribution(self.context, (9, )))
with self.assertRaises(ValueError):
source_da.distribute_as(Distribution(self.context, (3, 4)))
with self.assertRaises(ValueError):
source_da.distribute_as(Distribution(self.context, (3, 4000)))
def test_redist_unsupported_dist_types(self):
source_dist = Distribution(self.context, (10, 20), ('n', 'c'))
source_da = self.context.empty(source_dist)
with self.assertRaises(NotImplementedError):
source_da.distribute_as(source_da.distribution)
source_da = self.context.empty((10, 20), ('b', 'b'))
dest_dist = Distribution(self.context, (10, 20), ('b', 'c'))
with self.assertRaises(NotImplementedError):
source_da.distribute_as(dest_dist)
def test_redist_reshape_split_targets(self):
dist0 = Distribution(self.context, (40, ), ('b', ), (1, ), targets=[0])
dist1 = Distribution(self.context, (5, 8), ('b', 'n'), (2, ),
targets=[0, 1])
da_src = self.context.empty(dist0)
da_src.fill(-10)
da_dest = da_src.distribute_as(dist1)
expected = numpy.empty((5, 8))
expected.fill(-10)
assert_array_equal(da_dest.tondarray(), expected)
def test_tondarray(self):
distribution = Distribution(self.context, (3, 3))
dap = self.context.empty(distribution)
ndarr = numpy.arange(9).reshape(3, 3)
for (i, j), val in numpy.ndenumerate(ndarr):
dap[i, j] = ndarr[i, j]
numpy.testing.assert_array_equal(dap.tondarray(), ndarr)
def test_gh_435_regression_with_var(self):
dist = Distribution(self.context,
shape=(14, ),
dist=('b'),
targets=range(4))
darr = self.context.ones(dist)
darr.var()
def test_set_index_error(self):
distribution = Distribution(self.context, (10, ), dist={0: 'c'})
dap = self.context.empty(distribution)
with self.assertRaises(IndexError):
dap[11] = 55
with self.assertRaises(IndexError):
dap[-11] = 55
def test_grid_rank(self):
# regression test for issue #235
d = Distribution(self.context, (4, 4, 4),
dist=('b', 'n', 'b'),
grid_shape=(1, 1, 4))
a = self.context.empty(d)
self.assertEqual(a.grid_shape, (1, 1, 4))
def test_iteration(self):
size = 10
distribution = Distribution(self.context, (size, ), dist={0: 'c'})
dap = self.context.empty(distribution)
dap.fill(10)
for val in dap:
self.assertEqual(val, 10)
def fromndarray(self, arr, distribution=None):
"""Create a DistArray from an ndarray.
Parameters
----------
distribution : Distribution object, optional
If a Distribution object is not provided, one is created with
`Distribution(arr.shape)`.
Returns
-------
DistArray
A DistArray distributed as specified, using the values and dtype
from `arr`.
"""
if distribution is None:
distribution = Distribution(self, arr.shape)
out = self.empty(distribution, dtype=arr.dtype)
try:
out[...] = arr
except AttributeError:
# no slicing for a given map type; do it the slow way
for index, value in numpy.ndenumerate(arr):
out[index] = value
return out
def test_reduce_0D(self):
N = 10**5
dist = Distribution(self.context, (N, ))
new_dist = dist.reduce(axes=[0])
self.assertEqual(new_dist.dist, ())
self.assertSequenceEqual(new_dist.shape, ())
self.assertEqual(new_dist.grid_shape, ())
self.assertEqual(set(new_dist.targets), set(dist.targets[:1]))
def test___init__(self):
shape = (5, 5)
distribution = Distribution(self.context, shape, ('b', 'c'))
da = DistArray(distribution, dtype=int)
da.fill(42)
nda = numpy.empty(shape, dtype=int)
nda.fill(42)
assert_array_equal(da.tondarray(), nda)
def test_2D_bn(self):
nrows, ncols = 31, 53
cm = Distribution(self.context, (nrows, ncols), {0: 'b'}, (4, 1))
chunksize = (nrows // 4) + 1
for _ in range(100):
r, c = randrange(nrows), randrange(ncols)
rank = r // chunksize
self.assertSequenceEqual(cm.owning_ranks((r, c)), [rank])
def test_redist_identity(self):
source_dist = dest_dist = Distribution(self.context, (10, 10),
('b', 'b'), (1, 1),
targets=[0])
source_da = self.context.empty(source_dist, dtype=numpy.int32)
source_da.fill(-42)
dest_da = source_da.distribute_as(dest_dist)
assert_array_equal(source_da.tondarray(), dest_da.tondarray())
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)
def test_redist_reshape_three_dee(self):
three_dee_shape = (3, 5, 2)
two_dee_shape = (3 * 2, 5)
dist0 = Distribution(self.context, three_dee_shape, ('b', 'b', 'n'),
(2, 2))
dist1 = Distribution(self.context, two_dee_shape, ('b', 'b'), (2, 2))
da_src = self.context.empty(dist0)
da_src.fill(47)
da_dest = da_src.distribute_as(dist1)
self.assertTrue(da_dest.distribution.is_compatible(dist1))
expected = numpy.empty(two_dee_shape)
expected.fill(47)
assert_array_equal(da_dest.tondarray(), expected)
def local_view(larr, ddpr, dtype):
from distarray.localapi.maps import Distribution
if len(ddpr) == 0:
dim_data = ()
else:
dim_data = ddpr[larr.comm_rank]
ldist = Distribution(comm=larr.comm, dim_data=dim_data)
lview = larr.view(ldist, dtype=dtype)
return proxyize(lview)
def _local_create(comm, ddpr, dtype):
from distarray.localapi import empty
from distarray.localapi.maps import Distribution
if len(ddpr):
dim_data = ddpr[comm.Get_rank()]
else:
dim_data = ()
dist = Distribution(comm=comm, dim_data=dim_data)
return proxyize(empty(dist, dtype))
def get_view(arr, index, ddpr, comm):
from distarray.localapi.maps import Distribution
if len(ddpr) == 0:
dim_data = ()
else:
dim_data = ddpr[comm.Get_rank()]
local_distribution = Distribution(comm=comm, dim_data=dim_data)
result = arr.global_index.get_slice(index, local_distribution)
return proxyize(result)
def test_full_slice_with_int_2d(self):
d0 = Distribution(context=self.context, shape=(15, 20))
s = (slice(None), 4)
d1 = d0.slice(s)
self.assertEqual(len(d0.maps) - 1, len(d1.maps))
self.assertSequenceEqual(d1.dist, d0.dist[:-1])
self.assertEqual(d1.shape, (15, ))
def test_from_partial_slice_1d(self):
d0 = Distribution(context=self.context, shape=(15, ))
s = (slice(0, 3), )
d1 = d0.slice(s)
self.assertEqual(len(d0.maps), len(d1.maps))
self.assertSequenceEqual(d1.dist, d0.dist)
self.assertSequenceEqual(d1.targets, [0])
self.assertSequenceEqual(d1.shape, (3, ))
def create_local(local_call, ddpr, dtype, comm):
from distarray.localapi.maps import Distribution
if len(ddpr) == 0:
dim_data = ()
else:
dim_data = ddpr[comm.Get_rank()]
local_call = eval(local_call)
distribution = Distribution(comm=comm, dim_data=dim_data)
rval = local_call(distribution=distribution, dtype=dtype)
return proxyize(rval)
def test_from_full_slice_with_step_1d_0(self):
d0 = Distribution(context=self.context, shape=(15, ))
s = (slice(None, None, 2), )
d1 = d0.slice(s)
self.assertEqual(len(d0.maps), len(d1.maps))
self.assertSequenceEqual(d1.dist, d0.dist)
self.assertSequenceEqual(d1.targets, d0.targets)
self.assertEqual(d1.maps[0].bounds[0][0], d0.maps[0].bounds[0][0])
def test_from_partial_slice_2d(self):
d0 = Distribution(context=self.context, shape=(15, 20))
s = (slice(3, 7), 4)
d1 = d0.slice(s)
self.assertEqual(len(d0.maps) - 1, len(d1.maps))
self.assertSequenceEqual(d1.dist, d0.dist[:-1])
for m, expected in zip(d1.maps, ([(0, 1), (1, 4)], [(0, 1)])):
self.assertSequenceEqual(m.bounds, expected)
def test_block_redist_many_to_one(self):
source_dist = Distribution(self.context, (40, ), ('b', ), (2, ),
targets=[1, 2])
dest_dist = Distribution(self.context, (40, ), ('b', ), (1, ),
targets=[0])
plan = source_dist.get_redist_plan(dest_dist)
expected = [
{
'source_rank': 1,
'dest_rank': 0,
'indices': [(0, 20, 1)]
},
{
'source_rank': 2,
'dest_rank': 0,
'indices': [(20, 40, 1)]
},
]
self.assertEqual(plan, expected)
