From c3d615c6b8bc5ad0088b715f918f6f6f179361e2 Mon Sep 17 00:00:00 2001
From: Robert David Grant <rgrant@enthought.com>
Date: Fri, 2 May 2014 11:50:21 -0500
Subject: [PATCH 01/12] Tweak docstrings and blank lines.

---
 distarray/context.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/distarray/context.py b/distarray/context.py
index 74434670..5db771ef 100644
--- a/distarray/context.py
+++ b/distarray/context.py
@@ -3,9 +3,10 @@
 #  Copyright (C) 2008-2014, IPython Development Team and Enthought, Inc.
 #  Distributed under the terms of the BSD License.  See COPYING.rst.
 # ---------------------------------------------------------------------------
+
 """
-`Context` objects contain the information required for distarrays to
-communicate with localarrays.
+`Context` objects contain the information required for `DistArray`s to
+communicate with `LocalArray`s.
 """
 
 
@@ -26,6 +27,7 @@
 
 
 class Context(object):
+
     """
     Context objects manage the setup and communication of the worker processes
     for DistArray objects.  A DistArray object has a context, and contexts have
@@ -35,7 +37,6 @@ class Context(object):
     Typically there is just one context object that uses all processes,
     although it is possible to have more than one context with a different
     selection of engines.
-
     """
 
     _CLEANUP = None

From 2c490c9478c9b36eb7ff9f32efd9c8dc8be55049 Mon Sep 17 00:00:00 2001
From: Robert David Grant <rgrant@enthought.com>
Date: Fri, 2 May 2014 12:17:30 -0500
Subject: [PATCH 02/12] Group distribution-related keys in API.

---
 distarray/context.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/distarray/context.py b/distarray/context.py
index 5db771ef..4a9d06c0 100644
--- a/distarray/context.py
+++ b/distarray/context.py
@@ -217,21 +217,21 @@ def _create_local(self, local_call, shape, dist, grid_shape, dtype):
         return DistArray.from_localarrays(da_key, distribution=distribution,
                                           dtype=dtype)
 
-    def zeros(self, shape, dtype=float, dist=None, grid_shape=None):
+    def zeros(self, shape, dist=None, grid_shape=None, dtype=float):
         if dist is None:
             dist = {0: 'b'}
         return self._create_local(local_call='distarray.local.zeros',
                                   shape=shape, dist=dist,
                                   grid_shape=grid_shape, dtype=dtype)
 
-    def ones(self, shape, dtype=float, dist=None, grid_shape=None):
+    def ones(self, shape, dist=None, grid_shape=None, dtype=float):
         if dist is None:
             dist = {0: 'b'}
         return self._create_local(local_call='distarray.local.ones',
                                   shape=shape, dist=dist,
                                   grid_shape=grid_shape, dtype=dtype,)
 
-    def empty(self, shape, dtype=float, dist=None, grid_shape=None):
+    def empty(self, shape, dist=None, grid_shape=None, dtype=float):
         if dist is None:
             dist = {0: 'b'}
         return self._create_local(local_call='distarray.local.empty',
@@ -469,8 +469,8 @@ def fromndarray(self, arr, dist=None, grid_shape=None):
         """Convert an ndarray to a distarray."""
         if dist is None:
             dist = {0: 'b'}
-        out = self.empty(arr.shape, dtype=arr.dtype, dist=dist,
-                         grid_shape=grid_shape)
+        out = self.empty(arr.shape, dist=dist, grid_shape=grid_shape,
+                         dtype=arr.dtype)
         for index, value in numpy.ndenumerate(arr):
             out[index] = value
         return out

From 9a7c62b78b566eb39ad2842a27422cf04aea318f Mon Sep 17 00:00:00 2001
From: Robert David Grant <rgrant@enthought.com>
Date: Fri, 2 May 2014 12:32:18 -0500
Subject: [PATCH 03/12] Make load_npy take a Distribution

---
 distarray/context.py                   | 16 ++++------------
 distarray/tests/test_distributed_io.py | 16 ++++++++++------
 2 files changed, 14 insertions(+), 18 deletions(-)

diff --git a/distarray/context.py b/distarray/context.py
index 4a9d06c0..75bf61ab 100644
--- a/distarray/context.py
+++ b/distarray/context.py
@@ -387,7 +387,7 @@ def save_hdf5(self, filename, da, key='buffer', mode='a'):
             'distarray.local.save_hdf5(%s, %s, %s, %s)' % subs
         )
 
-    def load_npy(self, filename, dim_data_per_rank):
+    def load_npy(self, filename, distribution):
         """
         Load a DistArray from a dataset in a ``.npy`` file.
 
@@ -395,9 +395,7 @@ def load_npy(self, filename, dim_data_per_rank):
         ----------
         filename : str
             Filename to load.
-        dim_data_per_rank : sequence of tuples of dict
-            A "dim_data" data structure for every rank.  Described here:
-            https://github.com/enthought/distributed-array-protocol
+        distribution: Distribution object
 
         Returns
         -------
@@ -405,20 +403,14 @@ def load_npy(self, filename, dim_data_per_rank):
             A DistArray encapsulating the file loaded.
 
         """
-        if len(self.targets) != len(dim_data_per_rank):
-            errmsg = "`dim_data_per_rank` must contain a dim_data for every rank."
-            raise TypeError(errmsg)
-
         da_key = self._generate_key()
-        subs = ((da_key,) + self._key_and_push(filename, dim_data_per_rank) +
+        ddpr = distribution.get_dim_data_per_rank()
+        subs = ((da_key,) + self._key_and_push(filename, ddpr) +
                 (self._comm_key,) + (self._comm_key,))
 
         self._execute(
             '%s = distarray.local.load_npy(%s, %s[%s.Get_rank()], %s)' % subs
         )
-
-        distribution = Distribution.from_dim_data_per_rank(self,
-                                                           dim_data_per_rank)
         return DistArray.from_localarrays(da_key, distribution=distribution)
 
     def load_hdf5(self, filename, dim_data_per_rank, key='buffer'):
diff --git a/distarray/tests/test_distributed_io.py b/distarray/tests/test_distributed_io.py
index bfa967c7..b4a660ed 100644
--- a/distarray/tests/test_distributed_io.py
+++ b/distarray/tests/test_distributed_io.py
@@ -20,6 +20,7 @@
 from distarray.externals.six.moves import range
 
 from distarray.client import DistArray
+from distarray.client_map import Distribution
 from distarray.context import Context
 from distarray.testing import import_or_skip, temp_filepath
 
@@ -150,22 +151,25 @@ def tearDown(self):
         self.dac.close()
 
     def test_load_bn(self):
-        dim_data_per_rank = bn_test_data
-        da = self.dac.load_npy(self.output_path, dim_data_per_rank)
+        distribution = Distribution.from_dim_data_per_rank(self.dac,
+                                                           bn_test_data)
+        da = self.dac.load_npy(self.output_path, distribution)
         for i in range(da.shape[0]):
             for j in range(da.shape[1]):
                 self.assertEqual(da[i, j], self.expected[i, j])
 
     def test_load_nc(self):
-        dim_data_per_rank = nc_test_data
-        da = self.dac.load_npy(self.output_path, dim_data_per_rank)
+        distribution = Distribution.from_dim_data_per_rank(self.dac,
+                                                           nc_test_data)
+        da = self.dac.load_npy(self.output_path, distribution)
         for i in range(da.shape[0]):
             for j in range(da.shape[1]):
                 self.assertEqual(da[i, j], self.expected[i, j])
 
     def test_load_nu(self):
-        dim_data_per_rank = nu_test_data
-        da = self.dac.load_npy(self.output_path, dim_data_per_rank)
+        distribution = Distribution.from_dim_data_per_rank(self.dac,
+                                                           nu_test_data)
+        da = self.dac.load_npy(self.output_path, distribution)
         for i in range(da.shape[0]):
             for j in range(da.shape[1]):
                 self.assertEqual(da[i, j], self.expected[i, j])

From 148cf38468c0be0272e84127980bc6d2e138e334 Mon Sep 17 00:00:00 2001
From: Robert David Grant <rgrant@enthought.com>
Date: Fri, 2 May 2014 12:40:18 -0500
Subject: [PATCH 04/12] Make load_hdf5 take a distribution

---
 distarray/context.py                   | 17 ++++-------------
 distarray/tests/test_distributed_io.py | 12 +++++++++---
 2 files changed, 13 insertions(+), 16 deletions(-)

diff --git a/distarray/context.py b/distarray/context.py
index 75bf61ab..78af4cd3 100644
--- a/distarray/context.py
+++ b/distarray/context.py
@@ -413,7 +413,7 @@ def load_npy(self, filename, distribution):
         )
         return DistArray.from_localarrays(da_key, distribution=distribution)
 
-    def load_hdf5(self, filename, dim_data_per_rank, key='buffer'):
+    def load_hdf5(self, filename, distribution, key='buffer'):
         """
         Load a DistArray from a dataset in an ``.hdf5`` file.
 
@@ -421,9 +421,7 @@ def load_hdf5(self, filename, dim_data_per_rank, key='buffer'):
         ----------
         filename : str
             Filename to load.
-        dim_data_per_rank : sequence of tuples of dict
-            A "dim_data" data structure for every rank.  Described here:
-            https://github.com/enthought/distributed-array-protocol
+        distribution: Distribution object
         key : str, optional
             The identifier for the group to load the DistArray from (the
             default is 'buffer').
@@ -440,21 +438,14 @@ def load_hdf5(self, filename, dim_data_per_rank, key='buffer'):
             errmsg = "An MPI-enabled h5py must be available to use load_hdf5."
             raise ImportError(errmsg)
 
-        if len(self.targets) != len(dim_data_per_rank):
-            errmsg = "`dim_data_per_rank` must contain a dim_data for every rank."
-            raise TypeError(errmsg)
-
         da_key = self._generate_key()
-        subs = ((da_key,) + self._key_and_push(filename, dim_data_per_rank) +
+        ddpr = distribution.get_dim_data_per_rank()
+        subs = ((da_key,) + self._key_and_push(filename, ddpr) +
                 (self._comm_key,) + self._key_and_push(key) + (self._comm_key,))
 
         self._execute(
             '%s = distarray.local.load_hdf5(%s, %s[%s.Get_rank()], %s, %s)' % subs
         )
-
-        distribution = Distribution.from_dim_data_per_rank(self,
-                                                           dim_data_per_rank)
-
         return DistArray.from_localarrays(da_key, distribution=distribution)
 
     def fromndarray(self, arr, dist=None, grid_shape=None):
diff --git a/distarray/tests/test_distributed_io.py b/distarray/tests/test_distributed_io.py
index b4a660ed..9cadaaf8 100644
--- a/distarray/tests/test_distributed_io.py
+++ b/distarray/tests/test_distributed_io.py
@@ -251,17 +251,23 @@ def tearDown(self):
         self.dac.close()
 
     def test_load_bn(self):
-        da = self.dac.load_hdf5(self.output_path, bn_test_data, key="test")
+        distribution = Distribution.from_dim_data_per_rank(self.dac,
+                                                           bn_test_data)
+        da = self.dac.load_hdf5(self.output_path, distribution, key="test")
         for i, v in np.ndenumerate(self.expected):
             self.assertEqual(v, da[i])
 
     def test_load_nc(self):
-        da = self.dac.load_hdf5(self.output_path, nc_test_data, key="test")
+        distribution = Distribution.from_dim_data_per_rank(self.dac,
+                                                           nc_test_data)
+        da = self.dac.load_hdf5(self.output_path, distribution, key="test")
         for i, v in np.ndenumerate(self.expected):
             self.assertEqual(v, da[i])
 
     def test_load_nu(self):
-        da = self.dac.load_hdf5(self.output_path, nu_test_data, key="test")
+        distribution = Distribution.from_dim_data_per_rank(self.dac,
+                                                           nu_test_data)
+        da = self.dac.load_hdf5(self.output_path, distribution, key="test")
         for i, v in np.ndenumerate(self.expected):
             self.assertEqual(v, da[i])
 

From c6a86b881104e820840096fee8ad06aeb65250e3 Mon Sep 17 00:00:00 2001
From: Robert David Grant <rgrant@enthought.com>
Date: Fri, 2 May 2014 13:15:16 -0500
Subject: [PATCH 05/12] WIP: Change {empty,zeros,ones} to take distributions

instead of shape/dist/grid_shape
---
 distarray/context.py | 31 +++++++++----------------------
 1 file changed, 9 insertions(+), 22 deletions(-)

diff --git a/distarray/context.py b/distarray/context.py
index 78af4cd3..73ade0e4 100644
--- a/distarray/context.py
+++ b/distarray/context.py
@@ -200,14 +200,10 @@ def _push0(self, d):
     def _pull0(self, k):
         return self.view.pull(k, targets=self.targets[0], block=True)
 
-    def _create_local(self, local_call, shape, dist, grid_shape, dtype):
+    def _create_local(self, local_call, distribution, dtype):
         """Creates LocalArrays with the method named in `local_call`."""
         da_key = self._generate_key()
         comm_name = self._comm_key
-        distribution = Distribution.from_shape(context=self,
-                                               shape=shape,
-                                               dist=dist,
-                                               grid_shape=grid_shape)
         ddpr = distribution.get_dim_data_per_rank()
         ddpr_name, dtype_name =  self._key_and_push(ddpr, dtype)
         cmd = ('{da_key} = {local_call}(distarray.local.maps.Distribution('
@@ -217,26 +213,17 @@ def _create_local(self, local_call, shape, dist, grid_shape, dtype):
         return DistArray.from_localarrays(da_key, distribution=distribution,
                                           dtype=dtype)
 
-    def zeros(self, shape, dist=None, grid_shape=None, dtype=float):
-        if dist is None:
-            dist = {0: 'b'}
+    def empty(self, distribution, dtype=float):
+        return self._create_local(local_call='distarray.local.empty',
+                                  distribution=distribution, dtype=dtype)
+
+    def zeros(self, distribution, dtype=float):
         return self._create_local(local_call='distarray.local.zeros',
-                                  shape=shape, dist=dist,
-                                  grid_shape=grid_shape, dtype=dtype)
+                                  distribution=distribution, dtype=dtype)
 
-    def ones(self, shape, dist=None, grid_shape=None, dtype=float):
-        if dist is None:
-            dist = {0: 'b'}
+    def ones(self, distribution, dtype=float):
         return self._create_local(local_call='distarray.local.ones',
-                                  shape=shape, dist=dist,
-                                  grid_shape=grid_shape, dtype=dtype,)
-
-    def empty(self, shape, dist=None, grid_shape=None, dtype=float):
-        if dist is None:
-            dist = {0: 'b'}
-        return self._create_local(local_call='distarray.local.empty',
-                                  shape=shape, dist=dist,
-                                  grid_shape=grid_shape, dtype=dtype)
+                                  distribution=distribution, dtype=dtype,)
 
     def save_dnpy(self, name, da):
         """

From f1106d92092e834a20ba961211a38d6f7bc996fa Mon Sep 17 00:00:00 2001
From: Robert David Grant <rgrant@enthought.com>
Date: Fri, 2 May 2014 17:04:56 -0500
Subject: [PATCH 06/12] WIP: Change fromndarray to take a distribution object.

---
 distarray/context.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/distarray/context.py b/distarray/context.py
index 73ade0e4..f6b137e7 100644
--- a/distarray/context.py
+++ b/distarray/context.py
@@ -435,12 +435,11 @@ def load_hdf5(self, filename, distribution, key='buffer'):
         )
         return DistArray.from_localarrays(da_key, distribution=distribution)
 
-    def fromndarray(self, arr, dist=None, grid_shape=None):
+    def fromndarray(self, arr, distribution=None):
         """Convert an ndarray to a distarray."""
-        if dist is None:
-            dist = {0: 'b'}
-        out = self.empty(arr.shape, dist=dist, grid_shape=grid_shape,
-                         dtype=arr.dtype)
+        if distribution is None:
+            distribution = Distribution.from_shape(self, arr.shape)
+        out = self.empty(distribution, dtype=arr.dtype)
         for index, value in numpy.ndenumerate(arr):
             out[index] = value
         return out

From 98464331e090fd1f134f0963ba343631de01b06b Mon Sep 17 00:00:00 2001
From: Robert David Grant <rgrant@enthought.com>
Date: Fri, 2 May 2014 17:22:15 -0500
Subject: [PATCH 07/12] Fix a usage of `empty` in decorators.

---
 distarray/decorators.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/distarray/decorators.py b/distarray/decorators.py
index 3dd60150..1a898337 100644
--- a/distarray/decorators.py
+++ b/distarray/decorators.py
@@ -11,6 +11,7 @@
 import functools
 
 from distarray.client import DistArray
+from distarray.client_map import Distribution
 from distarray.context import Context
 from distarray.error import ContextError
 from distarray.utils import has_exactly_one
@@ -195,9 +196,7 @@ def __call__(self, *args, **kwargs):
         for arg in args:
             if isinstance(arg, DistArray):
                 # Create the output distarray.
-                out = context.empty(arg.shape, dtype=arg.dtype,
-                                         dist=arg.dist,
-                                         grid_shape=arg.grid_shape)
+                out = context.empty(arg.distribution, dtype=arg.dtype)
                 # parse args
                 args_str, kwargs_str = self.key_and_push_args(
                     args, kwargs, context=context,

From 0744baabe94c6e24cc8120c2a9f4e2b39071e80d Mon Sep 17 00:00:00 2001
From: Robert David Grant <rgrant@enthought.com>
Date: Fri, 2 May 2014 17:23:03 -0500
Subject: [PATCH 08/12] Fix tests.

---
 distarray/tests/test_client.py         | 50 +++++++++++++++++---------
 distarray/tests/test_context.py        | 21 +++++++----
 distarray/tests/test_decorators.py     | 30 ++++++++++------
 distarray/tests/test_distributed_io.py | 17 ++++++---
 4 files changed, 80 insertions(+), 38 deletions(-)

diff --git a/distarray/tests/test_client.py b/distarray/tests/test_client.py
index 38b0ff39..d8573988 100644
--- a/distarray/tests/test_client.py
+++ b/distarray/tests/test_client.py
@@ -33,7 +33,9 @@ def tearDown(self):
 
     def test_set_and_getitem_block_dist(self):
         size = 10
-        dap = self.dac.empty((size,), dist={0: 'b'})
+        distribution = Distribution.from_shape(self.dac, (size,),
+                                               dist={0: 'b'})
+        dap = self.dac.empty(distribution)
 
         for val in range(size):
             dap[val] = val
@@ -47,7 +49,9 @@ def test_set_and_getitem_block_dist(self):
 
     def test_set_and_getitem_nd_block_dist(self):
         size = 5
-        dap = self.dac.empty((size, size), dist={0: 'b', 1: 'b'})
+        distribution = Distribution.from_shape(self.dac, (size, size),
+                                               dist={0: 'b', 1: 'b'})
+        dap = self.dac.empty(distribution)
 
         for row in range(size):
             for col in range(size):
@@ -55,15 +59,16 @@ def test_set_and_getitem_nd_block_dist(self):
                 dap[row, col] = val
                 self.assertEqual(dap[row, col], val)
 
-        for row in range(1 ,size + 1):
+        for row in range(1, size + 1):
             for col in range(1, size + 1):
                 dap[-row, -col] = row + col
                 self.assertEqual(dap[-row, -col], row + col)
 
-
     def test_set_and_getitem_cyclic_dist(self):
         size = 10
-        dap = self.dac.empty((size,), dist={0: 'c'})
+        distribution = Distribution.from_shape(self.dac, (size,),
+                                               dist={0: 'c'})
+        dap = self.dac.empty(distribution)
 
         for val in range(size):
             dap[val] = val
@@ -74,14 +79,16 @@ def test_set_and_getitem_cyclic_dist(self):
             self.assertEqual(dap[-i], i)
 
     def test_get_index_error(self):
-        dap = self.dac.empty((10,), dist={0: 'c'})
+        distribution = Distribution.from_shape(self.dac, (10,), dist={0: 'c'})
+        dap = self.dac.empty(distribution)
         with self.assertRaises(IndexError):
             dap[11]
         with self.assertRaises(IndexError):
             dap[-11]
 
     def test_set_index_error(self):
-        dap = self.dac.empty((10,), dist={0: 'c'})
+        distribution = Distribution.from_shape(self.dac, (10,), dist={0: 'c'})
+        dap = self.dac.empty(distribution)
         with self.assertRaises(IndexError):
             dap[11] = 55
         with self.assertRaises(IndexError):
@@ -89,13 +96,16 @@ def test_set_index_error(self):
 
     def test_iteration(self):
         size = 10
-        dap = self.dac.empty((size,), dist={0: 'c'})
+        distribution = Distribution.from_shape(self.dac, (size,),
+                                               dist={0: 'c'})
+        dap = self.dac.empty(distribution)
         dap.fill(10)
         for val in dap:
             self.assertEqual(val, 10)
 
     def test_tondarray(self):
-        dap = self.dac.empty((3, 3))
+        distribution = Distribution.from_shape(self.dac, (3, 3))
+        dap = self.dac.empty(distribution)
         ndarr = numpy.arange(9).reshape(3, 3)
         for (i, j), val in numpy.ndenumerate(ndarr):
             dap[i, j] = ndarr[i, j]
@@ -103,7 +113,9 @@ def test_tondarray(self):
 
     def test_global_tolocal_bug(self):
         # gh-issue #154
-        dap = self.dac.zeros((3, 3), dist=('n', 'b'))
+        distribution = Distribution.from_shape(self.dac, (3, 3),
+                                               dist=('n', 'b'))
+        dap = self.dac.zeros(distribution)
         ndarr = numpy.zeros((3, 3))
         numpy.testing.assert_array_equal(dap.tondarray(), ndarr)
 
@@ -327,20 +339,22 @@ def test___init__(self):
 
     def test_zeros(self):
         shape = (16, 16)
-        zero_distarray = self.context.zeros(shape)
+        distribution = Distribution.from_shape(self.context, shape)
+        zero_distarray = self.context.zeros(distribution)
         zero_ndarray = numpy.zeros(shape)
         assert_array_equal(zero_distarray.tondarray(), zero_ndarray)
 
     def test_ones(self):
         shape = (16, 16)
-        one_distarray = self.context.ones(shape)
+        distribution = Distribution.from_shape(self.context, shape)
+        one_distarray = self.context.ones(distribution)
         one_ndarray = numpy.ones(shape)
         assert_array_equal(one_distarray.tondarray(), one_ndarray)
 
     def test_empty(self):
-        shape = (16, 16)
-        empty_distarray = self.context.empty(shape)
-        self.assertEqual(empty_distarray.shape, shape)
+        distribution = Distribution.from_shape(self.context, (16, 16))
+        empty_distarray = self.context.empty(distribution)
+        self.assertEqual(empty_distarray.shape, distribution.shape)
 
     def test_fromndarray(self):
         ndarr = numpy.arange(16).reshape(4, 4)
@@ -350,8 +364,10 @@ def test_fromndarray(self):
 
     def test_grid_rank(self):
         # regression test for issue #235
-        a = self.context.empty((4, 4, 4), dist=('b', 'n', 'b'),
-                               grid_shape=(1, 1, 4))
+        d = Distribution.from_shape(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_fromfunction(self):
diff --git a/distarray/tests/test_context.py b/distarray/tests/test_context.py
index 91c442e2..d50e3423 100644
--- a/distarray/tests/test_context.py
+++ b/distarray/tests/test_context.py
@@ -18,6 +18,7 @@
 import numpy
 
 from distarray import Context
+from distarray.client_map import Distribution
 from distarray.ipython_utils import IPythonClient
 from distarray.local import LocalArray
 
@@ -108,19 +109,27 @@ def tearDownClass(cls):
         cls.context.close()
 
     def test_1D(self):
-        a = self.context.empty((3,))
+        d = Distribution.from_shape(self.context, (3,))
+        a = self.context.empty(d)
         self.assertEqual(a.grid_shape, (3,))
 
     def test_2D(self):
-        a = self.context.empty((3, 3))
-        b = self.context.empty((3, 3), dist=('n', 'b'))
+        da = Distribution.from_shape(self.context, (3, 3))
+        a = self.context.empty(da)
+        db = Distribution.from_shape(self.context, (3, 3), dist=('n', 'b'))
+        b = self.context.empty(db)
         self.assertEqual(a.grid_shape, (3, 1))
         self.assertEqual(b.grid_shape, (1, 3))
 
     def test_3D(self):
-        a = self.context.empty((3, 3, 3))
-        b = self.context.empty((3, 3, 3), dist=('n', 'b', 'n'))
-        c = self.context.empty((3, 3, 3), dist=('n', 'n', 'b'))
+        da = Distribution.from_shape(self.context, (3, 3, 3))
+        a = self.context.empty(da)
+        db = Distribution.from_shape(self.context, (3, 3, 3),
+                                     dist=('n', 'b', 'n'))
+        b = self.context.empty(db)
+        dc = Distribution.from_shape(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))
diff --git a/distarray/tests/test_decorators.py b/distarray/tests/test_decorators.py
index b7357612..589ee286 100644
--- a/distarray/tests/test_decorators.py
+++ b/distarray/tests/test_decorators.py
@@ -18,6 +18,7 @@
 from numpy.testing import assert_array_equal
 
 from distarray.context import Context
+from distarray.client_map import Distribution
 from distarray.decorators import DecoratorBase, local, vectorize
 from distarray.error import ContextError
 
@@ -27,7 +28,8 @@ class TestDecoratorBase(TestCase):
     def test_determine_context(self):
         context = Context()
         context2 = Context()  # for cross Context checking
-        da = context.ones((2, 2))
+        distribution = Distribution.from_shape(context, (2, 2))
+        da = context.ones(distribution)
 
         def dummy_func(*args, **kwargs):
             fn = lambda x: x
@@ -45,7 +47,8 @@ def dummy_func(*args, **kwargs):
     def test_key_and_push_args(self):
         context = Context()
 
-        da = context.ones((2, 2))
+        distribution = Distribution.from_shape(context, (2, 2))
+        da = context.ones(distribution)
         db = da*2
 
         def dummy_func(*args, **kwargs):
@@ -141,7 +144,8 @@ def parameterless():
     @classmethod
     def setUpClass(cls):
         cls.context = Context()
-        cls.da = cls.context.empty((5, 5))
+        distribution = Distribution.from_shape(cls.context, (5, 5))
+        cls.da = cls.context.empty(distribution)
         cls.da.fill(2 * numpy.pi)
 
     @classmethod
@@ -152,7 +156,8 @@ def test_local(self):
         """Test the @local decorator"""
         context = Context()
 
-        da = context.empty((4, 4))
+        distribution = Distribution.from_shape(context, (4, 4))
+        da = context.empty(distribution)
         a = numpy.empty((4, 4))
 
         def fill_a(a):
@@ -175,8 +180,10 @@ def fill_da(da):
     def test_different_contexts(self):
         ctx1 = Context(targets=range(4))
         ctx2 = Context(targets=range(3))
-        da1 = ctx1.ones((10,))
-        da2 = ctx2.ones((10,))
+        distribution1 = Distribution.from_shape(ctx1, (10,))
+        distribution2 = Distribution.from_shape(ctx2, (10,))
+        da1 = ctx1.ones(distribution1)
+        da2 = ctx2.ones(distribution2)
         db1 = self.local_sin(da1)
         db2 = self.local_sin(da2)
         ndarr1 = db1.toarray()
@@ -209,13 +216,15 @@ def test_local_add_nums(self):
         self.assert_allclose(df, 2 * numpy.pi + 11 + 12 + 13)
 
     def test_local_add_distarrayproxies(self):
-        dg = self.context.empty((5, 5))
+        distribution = Distribution.from_shape(self.context, (5, 5))
+        dg = self.context.empty(distribution)
         dg.fill(33)
         dh = self.local_add_distarrayproxies(self.da, dg)
         self.assert_allclose(dh, 33 + 2 * numpy.pi)
 
     def test_local_add_mixed(self):
-        di = self.context.empty((5, 5))
+        distribution = Distribution.from_shape(self.context, (5, 5))
+        di = self.context.empty(distribution)
         di.fill(33)
         dj = self.local_add_mixed(self.da, 11, di, 12)
         self.assert_allclose(dj, 2 * numpy.pi + 11 + 33 + 12)
@@ -233,9 +242,10 @@ def test_local_add_kwargs(self):
         self.assert_allclose(dl, 2 * numpy.pi + 11 + 12)
 
     def test_local_add_supermix(self):
-        dm = self.context.empty((5, 5))
+        distribution = Distribution.from_shape(self.context, (5, 5))
+        dm = self.context.empty(distribution)
         dm.fill(22)
-        dn = self.context.empty((5, 5))
+        dn = self.context.empty(distribution)
         dn.fill(44)
         do = self.local_add_supermix(self.da, 11, dm, 33, dc=dn, num3=55)
         expected = 2 * numpy.pi + 11 + 22 + 33 + 44 + 55 + 66
diff --git a/distarray/tests/test_distributed_io.py b/distarray/tests/test_distributed_io.py
index 9cadaaf8..4c0fa79d 100644
--- a/distarray/tests/test_distributed_io.py
+++ b/distarray/tests/test_distributed_io.py
@@ -29,7 +29,8 @@ class TestDnpyFileIO(unittest.TestCase):
 
     def test_save_load_with_filenames(self):
         dac = Context()
-        da = dac.empty((100,), dist={0: 'b'})
+        distribution = Distribution.from_shape(dac, (100,), dist={0: 'b'})
+        da = dac.empty(distribution)
 
         output_paths = [temp_filepath() for target in dac.targets]
         try:
@@ -44,7 +45,8 @@ def test_save_load_with_filenames(self):
 
     def test_save_load_with_prefix(self):
         dac = Context()
-        da = dac.empty((100,), dist={0: 'b'})
+        distribution = Distribution.from_shape(dac, (100,), dist={0: 'b'})
+        da = dac.empty(distribution)
 
         output_path = temp_filepath()
         try:
@@ -189,7 +191,9 @@ def tearDown(self):
 
     def test_save_block(self):
         datalen = 33
-        da = self.dac.empty((datalen,), dist={0: 'b'})
+        distribution = Distribution.from_shape(self.dac, (datalen,),
+                                               dist={0: 'b'})
+        da = self.dac.empty(distribution)
         for i in range(datalen):
             da[i] = i
 
@@ -204,7 +208,8 @@ def test_save_3d(self):
         source = np.random.random(shape)
 
         dist = {0: 'b', 1: 'c', 2: 'n'}
-        da = self.dac.empty(shape, dist=dist)
+        distribution = Distribution.from_shape(self.dac, shape, dist=dist)
+        da = self.dac.empty(distribution)
 
         for i in range(shape[0]):
             for j in range(shape[1]):
@@ -218,7 +223,9 @@ def test_save_3d(self):
 
     def test_save_two_datasets(self):
         datalen = 33
-        da = self.dac.empty((datalen,), dist={0: 'b'})
+        distribution = Distribution.from_shape(self.dac, (datalen,),
+                                               dist={0: 'b'})
+        da = self.dac.empty(distribution)
 
         for i in range(datalen):
             da[i] = i

From f6b628692b99ccb6ee0bcf62a17ae941288edbec Mon Sep 17 00:00:00 2001
From: Robert David Grant <rgrant@enthought.com>
Date: Fri, 2 May 2014 17:38:07 -0500
Subject: [PATCH 09/12] Fix examples.

---
 examples/basic_demo.py                   | 11 +++++++----
 examples/julia-set/bench_dist.py         |  6 +++---
 examples/julia-set/julia_distarray.py    | 12 ++++++++----
 examples/julia-set/julia_ipython.py      |  5 ++---
 examples/make_distarray.py               |  8 +++++---
 examples/plot_distarray_block_axis_0.py  |  4 +++-
 examples/plot_distarray_block_axis_1.py  |  4 +++-
 examples/plot_distarray_block_block.py   |  4 +++-
 examples/plot_distarray_block_cyclic.py  |  4 +++-
 examples/plot_distarray_cyclic_cyclic.py |  4 +++-
 examples/plot_distarray_protocol.py      |  7 +++++--
 11 files changed, 45 insertions(+), 24 deletions(-)

diff --git a/examples/basic_demo.py b/examples/basic_demo.py
index cf52b644..bcb16be8 100644
--- a/examples/basic_demo.py
+++ b/examples/basic_demo.py
@@ -10,6 +10,7 @@
 import distarray
 from distarray.externals.six.moves import input
 from distarray import Context
+from distarray.client_map import Distribution
 from distarray.decorators import local
 from pprint import pprint
 
@@ -42,8 +43,8 @@ def global_sum(da):
     global_sum = da.distribution.comm.allreduce(local_sum, None, op=MPI.SUM)
 
     new_arr = numpy.array([global_sum])
-    distribution = distarray.local.maps.Distribution.from_shape((1,))
-    new_distarray = distarray.local.LocalArray(distribution, buf=new_arr)
+    distribution = Distribution.from_shape((1,))
+    new_distarray = LocalArray(distribution, buf=new_arr)
     return new_distarray
 
 
@@ -53,8 +54,10 @@ def global_sum(da):
 
     print()
     input("Basic creation:")
-    dap_b = context.empty((arr_len,), dist={0: 'b'})
-    dap_c = context.empty((arr_len,), dist={0: 'c'})
+    dist_b = Distribution.from_shape(context, (arr_len,), dist={0: 'b'})
+    dap_b = context.empty(dist_b)
+    dist_c = Distribution.from_shape(context, (arr_len,), dist={0: 'c'})
+    dap_c = context.empty(dist_c)
     print("dap_b is a ", type(dap_b))
     print("dap_c is a ", type(dap_c))
 
diff --git a/examples/julia-set/bench_dist.py b/examples/julia-set/bench_dist.py
index b936e927..8a2c35a5 100644
--- a/examples/julia-set/bench_dist.py
+++ b/examples/julia-set/bench_dist.py
@@ -17,15 +17,15 @@
 from IPython.parallel import Client
 from matplotlib import pyplot
 
-from distarray import Context
+from distarray import Context, Distribution
 from distarray.decorators import local, vectorize
 
 
 # Make an empty distributed array
 def make_empty_da(resolution, dist, context):
     """Create the arr we will build the fractal with."""
-    out = context.empty((resolution[0], resolution[1]),
-                        dist=dist, dtype=complex)
+    distribution = Distribution.from_shape(context, (resolution[0], resolution[1]), dist=dist)
+    out = context.empty(distribution, dtype=complex)
     return out
 
 
diff --git a/examples/julia-set/julia_distarray.py b/examples/julia-set/julia_distarray.py
index 17b020ed..9522def3 100644
--- a/examples/julia-set/julia_distarray.py
+++ b/examples/julia-set/julia_distarray.py
@@ -15,6 +15,7 @@
 from matplotlib import pyplot
 
 from distarray import Context
+from distarray.client_map import Distribution
 from distarray.decorators import local, vectorize
 
 context = Context()
@@ -25,8 +26,11 @@
 # Make an empty distributed array
 def make_empty_da(resolution, dist):
     """Create the arr we will build the fractal with."""
-    out = context.empty((resolution[0], resolution[1]),
-                        dist=dist, dtype=complex)
+
+    distribution = Distribution.from_shape(context,
+                                           (resolution[0], resolution[1]),
+                                           dist=dist)
+    out = context.empty(distribution, dtype=complex)
     return out
 
 
@@ -37,8 +41,8 @@ def draw_coord(arr, re_ax, im_ax, resolution):
     """Draw the complex coordinate plane"""
     re_step = float(re_ax[1] - re_ax[0]) / resolution[0]
     im_step = float(im_ax[1] - im_ax[0]) / resolution[1]
-    for i in arr.maps[0].global_iter:
-        for j in arr.maps[1].global_iter:
+    for i in arr.distribution[0].global_iter:
+        for j in arr.distribution[1].global_iter:
             arr.global_index[i, j] = complex(re_ax[0] + re_step*i,
                                              im_ax[0] + im_step*j)
     return arr
diff --git a/examples/julia-set/julia_ipython.py b/examples/julia-set/julia_ipython.py
index 1d850ab8..1b90ab0f 100644
--- a/examples/julia-set/julia_ipython.py
+++ b/examples/julia-set/julia_ipython.py
@@ -73,15 +73,14 @@ def calc_julia(dview, fn, c, z_max, n_max, re_ax, im_ax):
 if __name__ == '__main__':
     import sys
 
+    import numpy
     from IPython.parallel import Client
     from matplotlib import pyplot
 
     client = Client()
     dview = client[:]
     dview.block = True
-
-    with dview.sync_imports():
-        import numpy
+    dview.execute("import numpy")
 
     fn = lambda z, c: z*z + c
     c = numpy.complex(float(sys.argv[1]), float(sys.argv[2]))
diff --git a/examples/make_distarray.py b/examples/make_distarray.py
index 99038ae8..b06756c1 100644
--- a/examples/make_distarray.py
+++ b/examples/make_distarray.py
@@ -8,8 +8,10 @@
 Create a distarray.
 """
 
-import distarray
+from distarray.context import Context
+from distarray.client_map import Distribution
 
 
-c = distarray.Context()
-a = c.zeros((10, 10, 10), dtype='int32', dist=('b', 'n', 'c'))
+c = Context()
+d = Distribution.from_shape(c, (10, 10, 10), dist=('b', 'n', 'c'))
+a = c.zeros(d, dtype='int32')
diff --git a/examples/plot_distarray_block_axis_0.py b/examples/plot_distarray_block_axis_0.py
index 6e8c6b75..182c274d 100644
--- a/examples/plot_distarray_block_axis_0.py
+++ b/examples/plot_distarray_block_axis_0.py
@@ -13,10 +13,12 @@
 
 import distarray
 from distarray import plotting
+from distarray.client_map import Distribution
 
 
 c = distarray.Context()
-a = c.zeros((64, 64))
+d = Distribution.from_shape(c, (64, 64))
+a = c.zeros(d)
 process_coords = [(0, 0), (1, 0), (2, 0), (3, 0)]
 plotting.plot_array_distribution(a, process_coords, cell_label=False,
                                  legend=True)
diff --git a/examples/plot_distarray_block_axis_1.py b/examples/plot_distarray_block_axis_1.py
index f98f4591..77630bff 100644
--- a/examples/plot_distarray_block_axis_1.py
+++ b/examples/plot_distarray_block_axis_1.py
@@ -12,10 +12,12 @@
 
 import distarray
 from distarray import plotting
+from distarray.client_map import Distribution
 
 
 c = distarray.Context()
-a = c.zeros((64, 64), dtype='int32', dist=('n', 'b'))
+d = Distribution.from_shape(c, (64, 64), dist=('n', 'b'))
+a = c.zeros(d, dtype='int32')
 process_coords = [(0, 0), (1, 0), (2, 0), (3, 0)]
 plotting.plot_array_distribution(a, process_coords, cell_label=False,
                                  legend=True)
diff --git a/examples/plot_distarray_block_block.py b/examples/plot_distarray_block_block.py
index a4c14d5f..e1aad058 100644
--- a/examples/plot_distarray_block_block.py
+++ b/examples/plot_distarray_block_block.py
@@ -12,10 +12,12 @@
 
 import distarray
 from distarray import plotting
+from distarray.client_map import Distribution
 
 
 c = distarray.Context()
-a = c.zeros((64, 64), dtype='int32', dist=('b', 'b'))
+d = Distribution.from_shape(c, (64, 64), dist=('b', 'b'))
+a = c.zeros(d, dtype='int32')
 process_coords = [(0, 0), (0, 1), (1, 0), (1, 1)]
 plotting.plot_array_distribution(a, process_coords, cell_label=False,
                                  legend=True)
diff --git a/examples/plot_distarray_block_cyclic.py b/examples/plot_distarray_block_cyclic.py
index 63ad3344..e305f083 100644
--- a/examples/plot_distarray_block_cyclic.py
+++ b/examples/plot_distarray_block_cyclic.py
@@ -12,10 +12,12 @@
 
 import distarray
 from distarray import plotting
+from distarray.client_map import Distribution
 
 
 c = distarray.Context()
-a = c.zeros((64, 64), dtype='int32', dist=('b', 'c'))
+d = Distribution.from_shape(c, (64, 64), dist=('b', 'c'))
+a = c.zeros(d, dtype='int32')
 process_coords = [(0, 0), (0, 1), (1, 0), (1, 1)]
 plotting.plot_array_distribution(a, process_coords, cell_label=False,
                                  legend=True)
diff --git a/examples/plot_distarray_cyclic_cyclic.py b/examples/plot_distarray_cyclic_cyclic.py
index 651c4a19..989141c4 100644
--- a/examples/plot_distarray_cyclic_cyclic.py
+++ b/examples/plot_distarray_cyclic_cyclic.py
@@ -12,10 +12,12 @@
 
 import distarray
 from distarray import plotting
+from distarray.client_map import Distribution
 
 
 c = distarray.Context()
-a = c.zeros((64, 64), dtype='int32', dist=('c', 'c'))
+d = Distribution.from_shape(c, (64, 64), dist=('c', 'c'))
+a = c.zeros(d, dtype='int32')
 process_coords = [(0, 0), (0, 1), (1, 0), (1, 1)]
 plotting.plot_array_distribution(a, process_coords, cell_label=False,
                                  legend=True)
diff --git a/examples/plot_distarray_protocol.py b/examples/plot_distarray_protocol.py
index 52ffc066..f8cd5338 100644
--- a/examples/plot_distarray_protocol.py
+++ b/examples/plot_distarray_protocol.py
@@ -31,6 +31,7 @@
 
 import distarray
 from distarray import plotting
+from distarray.client_map import Distribution
 
 
 def print_array_documentation(context,
@@ -302,11 +303,13 @@ def shape_text(shape):
 
     # Create array, either from dist or dimdata.
     if dist is not None:
-        array = context.empty(shape, dist=dist, grid_shape=grid_shape)
+        distribution = Distribution.from_shape(context, shape, dist=dist,
+                                               grid_shape=grid_shape)
     elif dimdata is not None:
-        array = context.from_global_dim_data(dimdata)
+        distribution = Distribution(context, dimdata)
     else:
         raise ValueError('Must provide either dist or dimdata.')
+    array = context.empty(distribution)
 
     # Fill the array. This is slow but not a real problem here.
     value = 0.0

From 34a0254db5286b559a80c2eff27d4c2072045492 Mon Sep 17 00:00:00 2001
From: Robert David Grant <rgrant@enthought.com>
Date: Fri, 2 May 2014 17:58:21 -0500
Subject: [PATCH 10/12] Add simple docstrings.

---
 distarray/context.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/distarray/context.py b/distarray/context.py
index f6b137e7..c758ae56 100644
--- a/distarray/context.py
+++ b/distarray/context.py
@@ -214,14 +214,17 @@ def _create_local(self, local_call, distribution, dtype):
                                           dtype=dtype)
 
     def empty(self, distribution, dtype=float):
+        """Create an empty Distarray."""
         return self._create_local(local_call='distarray.local.empty',
                                   distribution=distribution, dtype=dtype)
 
     def zeros(self, distribution, dtype=float):
+        """Create a Distarray filled with zeros."""
         return self._create_local(local_call='distarray.local.zeros',
                                   distribution=distribution, dtype=dtype)
 
     def ones(self, distribution, dtype=float):
+        """Create a Distarray filled with ones."""
         return self._create_local(local_call='distarray.local.ones',
                                   distribution=distribution, dtype=dtype,)
 
@@ -436,7 +439,7 @@ def load_hdf5(self, filename, distribution, key='buffer'):
         return DistArray.from_localarrays(da_key, distribution=distribution)
 
     def fromndarray(self, arr, distribution=None):
-        """Convert an ndarray to a distarray."""
+        """Create a DistArray from an ndarray."""
         if distribution is None:
             distribution = Distribution.from_shape(self, arr.shape)
         out = self.empty(distribution, dtype=arr.dtype)

From f1b27f20269164642b6dd75396245d27e8c265b8 Mon Sep 17 00:00:00 2001
From: Robert David Grant <rgrant@enthought.com>
Date: Fri, 2 May 2014 18:07:26 -0500
Subject: [PATCH 11/12] Add better docstrings.

---
 distarray/context.py | 54 ++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 50 insertions(+), 4 deletions(-)

diff --git a/distarray/context.py b/distarray/context.py
index c758ae56..b61385a5 100644
--- a/distarray/context.py
+++ b/distarray/context.py
@@ -214,17 +214,50 @@ def _create_local(self, local_call, distribution, dtype):
                                           dtype=dtype)
 
     def empty(self, distribution, dtype=float):
-        """Create an empty Distarray."""
+        """Create an empty Distarray.
+
+        Parameters
+        ----------
+        distribution : Distribution object
+        dtype : NumPy dtype, optional (default float)
+
+        Returns
+        -------
+        DistArray
+            A DistArray distributed as specified, with uninitialized values.
+        """
         return self._create_local(local_call='distarray.local.empty',
                                   distribution=distribution, dtype=dtype)
 
     def zeros(self, distribution, dtype=float):
-        """Create a Distarray filled with zeros."""
+        """Create a Distarray filled with zeros.
+
+        Parameters
+        ----------
+        distribution : Distribution object
+        dtype : NumPy dtype, optional (default float)
+
+        Returns
+        -------
+        DistArray
+            A DistArray distributed as specified, filled with zeros.
+        """
         return self._create_local(local_call='distarray.local.zeros',
                                   distribution=distribution, dtype=dtype)
 
     def ones(self, distribution, dtype=float):
-        """Create a Distarray filled with ones."""
+        """Create a Distarray filled with ones.
+
+        Parameters
+        ----------
+        distribution : Distribution object
+        dtype : NumPy dtype, optional (default float)
+
+        Returns
+        -------
+        DistArray
+            A DistArray distributed as specified, filled with ones.
+        """
         return self._create_local(local_call='distarray.local.ones',
                                   distribution=distribution, dtype=dtype,)
 
@@ -439,7 +472,20 @@ def load_hdf5(self, filename, distribution, key='buffer'):
         return DistArray.from_localarrays(da_key, distribution=distribution)
 
     def fromndarray(self, arr, distribution=None):
-        """Create a DistArray from an ndarray."""
+        """Create a DistArray from an ndarray.
+
+        Parameters
+        ----------
+        distribution : Distribution object, optional
+            If a Distribution object is not provided, one is created with
+            `Distribution.from_shape(arr.shape)`.
+
+        Returns
+        -------
+        DistArray
+            A DistArray distributed as specified, using the values and dtype
+            from `arr`.
+        """
         if distribution is None:
             distribution = Distribution.from_shape(self, arr.shape)
         out = self.empty(distribution, dtype=arr.dtype)

From 9a0b1005b0e4161568730bddec40369d31016394 Mon Sep 17 00:00:00 2001
From: Robert David Grant <rgrant@enthought.com>
Date: Sat, 3 May 2014 23:41:43 -0500
Subject: [PATCH 12/12] Modify client-side random functions to take a
 distribution.

Also, fix the tests.

All of these functions have been modified to take a distribution as a required argument, whereas before, some of them could be called without shape/dist/grid_shape parameters to return a single scalar.  This didn't seem to be an important use-case (one could just use the appropriate numpy function for that), so I omitted it.
---
 distarray/random.py            | 103 ++++++---------------------------
 distarray/tests/test_random.py |  40 ++++++++-----
 2 files changed, 45 insertions(+), 98 deletions(-)

diff --git a/distarray/random.py b/distarray/random.py
index 2eb101f8..652f9838 100644
--- a/distarray/random.py
+++ b/distarray/random.py
@@ -40,11 +40,8 @@ def seed(self, seed=None):
             self.context._comm_key)
         self.context._execute(cmd)
 
-    def rand(self, size=None, dist=None, grid_shape=None):
-        """
-        rand(size=(d0, d1, ..., dn))
-
-        Random values in a given shape.
+    def rand(self, distribution):
+        """Random values over a given distribution.
 
         Create a distarray of the given shape and propagate it with
         random samples from a uniform distribution
@@ -52,28 +49,15 @@ def rand(self, size=None, dist=None, grid_shape=None):
 
         Parameters
         ----------
-        size : tuple of ints
-            The dimensions of the returned array, should all be positive.
-            If no argument is given a single Python float is returned.
-        dist : dist dictionary
-            Dictionary describing how to distribute the array along each axis.
-        grid_shape : tuple
-            Tuple describing the processor grid topology.
+        distribution : Distribution object
 
         Returns
         -------
-        out : distarray, shape ``(d0, d1, ..., dn)``
+        out : DistArray
             Random values.
 
         """
-        if dist is None:
-            dist = {0: 'b'}
         da_key = self.context._generate_key()
-
-        distribution = Distribution.from_shape(context=self.context,
-                                               shape=size,
-                                               dist=dist,
-                                               grid_shape=grid_shape)
         ddpr = distribution.get_dim_data_per_rank()
         ddpr_name = self.context._key_and_push(ddpr)[0]
         comm_name = self.context._comm_key
@@ -84,12 +68,8 @@ def rand(self, size=None, dist=None, grid_shape=None):
             'comm={comm_name}))'.format(**locals()))
         return DistArray.from_localarrays(da_key, distribution=distribution)
 
-    def normal(self, loc=0.0, scale=1.0, size=None, dist=None,
-               grid_shape=None):
-        """
-        normal(loc=0.0, scale=1.0, size=None, dist={0: 'b'}, grid_shape=None)
-
-        Draw random samples from a normal (Gaussian) distribution.
+    def normal(self, distribution, loc=0.0, scale=1.0):
+        """Draw random samples from a normal (Gaussian) distribution.
 
         The probability density function of the normal distribution, first
         derived by De Moivre and 200 years later by both Gauss and Laplace
@@ -107,13 +87,7 @@ def normal(self, loc=0.0, scale=1.0, size=None, dist=None,
             Mean ("centre") of the distribution.
         scale : float
             Standard deviation (spread or "width") of the distribution.
-        size : tuple of ints
-            Output shape.  If the given shape is, e.g., ``(m, n, k)``, then
-            ``m * n * k`` samples are drawn.
-        dist : dist dictionary
-            Dictionary describing how to distribute the array along each axis.
-        grid_shape : tuple
-            Tuple describing the processor grid topology.
+        distribution : Distribution object
 
         Notes
         -----
@@ -141,14 +115,7 @@ def normal(self, loc=0.0, scale=1.0, size=None, dist=None,
                pp. 51, 51, 125.
 
         """
-        if dist is None:
-            dist = {0: 'b'}
         da_key = self.context._generate_key()
-
-        distribution = Distribution.from_shape(context=self.context,
-                                               shape=size,
-                                               dist=dist,
-                                               grid_shape=grid_shape)
         ddpr = distribution.get_dim_data_per_rank()
         loc_name, scale_name, ddpr_name = \
             self.context._key_and_push(loc, scale, ddpr)
@@ -161,11 +128,8 @@ def normal(self, loc=0.0, scale=1.0, size=None, dist=None,
             'comm={comm_name}))'.format(**locals()))
         return DistArray.from_localarrays(da_key, distribution=distribution)
 
-    def randint(self, low, high=None, size=None, dist=None, grid_shape=None):
-        """
-        randint(low, high=None, size=None)
-
-        Return random integers from `low` (inclusive) to `high` (exclusive).
+    def randint(self, distribution, low, high=None):
+        """Return random integers from `low` (inclusive) to `high` (exclusive).
 
         Return random integers from the "discrete uniform" distribution in the
         "half-open" interval [`low`, `high`). If `high` is None (the default),
@@ -173,6 +137,7 @@ def randint(self, low, high=None, size=None, dist=None, grid_shape=None):
 
         Parameters
         ----------
+        distribution : Distribution object
         low : int
             Lowest (signed) integer to be drawn from the distribution (unless
             ``high=None``, in which case this parameter is the *highest* such
@@ -180,29 +145,14 @@ def randint(self, low, high=None, size=None, dist=None, grid_shape=None):
         high : int, optional
             if provided, one above the largest (signed) integer to be drawn
             from the distribution (see above for behavior if ``high=None``).
-        size : int or tuple of ints, optional
-            Output shape. Default is None, in which case a single int is
-            returned.
-        dist : dist dictionary
-            Dictionary describing how to distribute the array along each axis.
-        grid_shape : tuple
-            Tuple describing the processor grid topology.
 
         Returns
         -------
-        out : distarray of ints
-            `size`-shaped distarray of random integers from the appropriate
-            distribution, or a single such random int if `size` not provided.
+        out : DistArray of ints
+            DistArray of random integers from the appropriate distribution.
 
         """
-        if dist is None:
-            dist = {0: 'b'}
         da_key = self.context._generate_key()
-
-        distribution = Distribution.from_shape(context=self.context,
-                                               shape=size,
-                                               dist=dist,
-                                               grid_shape=grid_shape)
         ddpr = distribution.get_dim_data_per_rank()
         low_name, high_name, ddpr_name = \
             self.context._key_and_push(low, high, ddpr)
@@ -215,37 +165,20 @@ def randint(self, low, high=None, size=None, dist=None, grid_shape=None):
             'comm={comm_name}))'.format(**locals()))
         return DistArray.from_localarrays(da_key, distribution=distribution)
 
-    def randn(self, size=None, dist=None, grid_shape=None):
-        """
-        randn(size=(d0, d1, ..., dn))
-
-        Return samples from the "standard normal" distribution.
+    def randn(self, distribution):
+        """Return samples from the "standard normal" distribution.
 
         Parameters
         ----------
-        size : tuple of ints
-            The dimensions of the returned array, should be all positive.
-            If no argument is given a single Python float is returned.
-        dist : dist dictionary
-            Dictionary describing how to distribute the array along each axis.
-        grid_shape : tuple
-            Tuple describing the processor grid topology.
+        distribution : Distribution object
 
         Returns
         -------
-        out : distarray
-            A ``(d0, d1, ..., dn)``-shaped distarray of floating-point samples
-            from the standard normal distribution.
-
+        out : DistArray
+            A DistArray of floating-point samples from the standard normal
+            distribution.
         """
-        if dist is None:
-            dist = {0: 'b'}
         da_key = self.context._generate_key()
-
-        distribution = Distribution.from_shape(context=self.context,
-                                               shape=size,
-                                               dist=dist,
-                                               grid_shape=grid_shape)
         ddpr = distribution.get_dim_data_per_rank()
         ddpr_name = self.context._key_and_push(ddpr)[0]
         comm_name = self.context._comm_key
diff --git a/distarray/tests/test_random.py b/distarray/tests/test_random.py
index 4a94e737..0a2c2c4f 100644
--- a/distarray/tests/test_random.py
+++ b/distarray/tests/test_random.py
@@ -9,6 +9,7 @@
 import unittest
 
 from distarray.context import Context
+from distarray.client_map import Distribution
 from distarray.random import Random
 
 
@@ -29,36 +30,45 @@ def tearDownClass(cls):
 
     def test_rand(self):
         shape = (3, 4)
-        a = self.random.rand(shape)
+        distribution = Distribution.from_shape(context=self.context,
+                                               shape=shape)
+        a = self.random.rand(distribution)
         self.assertEqual(a.shape, shape)
 
     def test_normal(self):
-        size = (3, 4)  # aka shape
-        a = self.random.normal(size=size)
-        self.assertEqual(a.shape, size)
+        shape = (3, 4)  # aka shape
+        distribution = Distribution.from_shape(context=self.context,
+                                               shape=shape)
+        a = self.random.normal(distribution)
+        self.assertEqual(a.shape, shape)
 
     def test_randint(self):
         low = 1
-        size = (3, 4)  # aka shape
-        a = self.random.randint(low, size=size)
-        self.assertEqual(a.shape, size)
+        shape = (3, 4)  # aka shape
+        distribution = Distribution.from_shape(context=self.context,
+                                               shape=shape)
+        a = self.random.randint(distribution, low=low)
+        self.assertEqual(a.shape, shape)
 
     def test_randn(self):
         shape = (3, 4)
-        a = self.random.randn(shape)
+        distribution = Distribution.from_shape(context=self.context,
+                                               shape=shape)
+        a = self.random.randn(distribution)
         self.assertEqual(a.shape, shape)
 
     def test_seed_same(self):
         """ Test that the same seed generates the same sequence. """
         shape = (8, 6)
+        distribution = Distribution.from_shape(self.context, shape)
         seed = 0xfeedbeef
         # Seed and get some random numbers.
         self.random.seed(seed)
-        a = self.random.rand(shape)
+        a = self.random.rand(distribution)
         aa = a.toarray()
         # Seed again and get more random numbers.
         self.random.seed(seed)
-        b = self.random.rand(shape)
+        b = self.random.rand(distribution)
         bb = b.toarray()
         # For an explicit seed, these should match exactly.
         self.assertTrue((aa == bb).all())
@@ -68,11 +78,12 @@ def test_seed_none(self):
         shape = (8, 6)
         # Seed and get some random numbers.
         self.random.seed(None)
-        a = self.random.rand(shape)
+        d = Distribution.from_shape(self.context, shape)
+        a = self.random.rand(d)
         aa = a.toarray()
         # Seed again and get more random numbers.
         self.random.seed(None)
-        b = self.random.rand(shape)
+        b = self.random.rand(d)
         bb = b.toarray()
         # For seed=None, these should *not* match.
         self.assertFalse((aa == bb).all())
@@ -86,7 +97,10 @@ def get_rand_array_per_engine(context, num_cols):
             """
             num_engines = len(context.targets)
             shape = (num_engines, num_cols)
-            darr = self.random.rand(size=shape, dist={0: 'b', 1: 'n'})
+            distribution = Distribution.from_shape(self.context,
+                                                   shape=shape,
+                                                   dist={0: 'b', 1: 'n'})
+            darr = self.random.rand(distribution)
             return darr
 
         # Seed generators.