Run flake8 in Travis and make code PEP8 compliant. (#387)

python/ray/experimental/array/distributed/__init__.py

@@ -4,4 +4,10 @@ from __future__ import print_function
 
 from . import random
 from . import linalg
-from .core import *
+from .core import (BLOCK_SIZE, DistArray, assemble, zeros, ones, copy, eye,
+                   triu, tril, blockwise_dot, dot, transpose, add, subtract,
+                   numpy_to_dist, subblocks)
+
+__all__ = ["random", "linalg", "BLOCK_SIZE", "DistArray", "assemble", "zeros",
+           "ones", "copy", "eye", "triu", "tril", "blockwise_dot", "dot",
+           "transpose", "add", "subtract", "numpy_to_dist", "subblocks"]

python/ray/experimental/array/distributed/core.py

@@ -6,30 +6,38 @@ import numpy as np
 import ray.experimental.array.remote as ra
 import ray
 
-__all__ = ["BLOCK_SIZE", "DistArray", "assemble", "zeros", "ones", "copy",
-           "eye", "triu", "tril", "blockwise_dot", "dot", "transpose", "add", "subtract", "numpy_to_dist", "subblocks"]
-
 BLOCK_SIZE = 10
 
+
 class DistArray(object):
   def __init__(self, shape, objectids=None):
     self.shape = shape
     self.ndim = len(shape)
     self.num_blocks = [int(np.ceil(1.0 * a / BLOCK_SIZE)) for a in self.shape]
-    self.objectids = objectids if objectids is not None else np.empty(self.num_blocks, dtype=object)
+    if objectids is not None:
+      self.objectids = objectids
+    else:
+      self.objectids = np.empty(self.num_blocks, dtype=object)
     if self.num_blocks != list(self.objectids.shape):
-      raise Exception("The fields `num_blocks` and `objectids` are inconsistent, `num_blocks` is {} and `objectids` has shape {}".format(self.num_blocks, list(self.objectids.shape)))
+      raise Exception("The fields `num_blocks` and `objectids` are "
+                      "inconsistent, `num_blocks` is {} and `objectids` has "
+                      "shape {}".format(self.num_blocks,
+                                        list(self.objectids.shape)))
 
   @staticmethod
   def compute_block_lower(index, shape):
     if len(index) != len(shape):
-      raise Exception("The fields `index` and `shape` must have the same length, but `index` is {} and `shape` is {}.".format(index, shape))
+      raise Exception("The fields `index` and `shape` must have the same "
+                      "length, but `index` is {} and `shape` is "
+                      "{}.".format(index, shape))
     return [elem * BLOCK_SIZE for elem in index]
 
   @staticmethod
   def compute_block_upper(index, shape):
     if len(index) != len(shape):
-      raise Exception("The fields `index` and `shape` must have the same length, but `index` is {} and `shape` is {}.".format(index, shape))
+      raise Exception("The fields `index` and `shape` must have the same "
+                      "length, but `index` is {} and `shape` is "
+                      "{}.".format(index, shape))
     upper = []
     for i in range(len(shape)):
       upper.append(min((index[i] + 1) * BLOCK_SIZE, shape[i]))
@@ -46,59 +54,73 @@ class DistArray(object):
     return [int(np.ceil(1.0 * a / BLOCK_SIZE)) for a in shape]
 
   def assemble(self):
-    """Assemble an array on this node from a distributed array of object IDs."""
+    """Assemble an array from a distributed array of object IDs."""
     first_block = ray.get(self.objectids[(0,) * self.ndim])
     dtype = first_block.dtype
     result = np.zeros(self.shape, dtype=dtype)
     for index in np.ndindex(*self.num_blocks):
       lower = DistArray.compute_block_lower(index, self.shape)
       upper = DistArray.compute_block_upper(index, self.shape)
-      result[[slice(l, u) for (l, u) in zip(lower, upper)]] = ray.get(self.objectids[index])
+      result[[slice(l, u) for (l, u) in zip(lower, upper)]] = ray.get(
+          self.objectids[index])
     return result
 
   def __getitem__(self, sliced):
-    # TODO(rkn): fix this, this is just a placeholder that should work but is inefficient
+    # TODO(rkn): Fix this, this is just a placeholder that should work but is
+    # inefficient.
     a = self.assemble()
     return a[sliced]
 
+
 # Register the DistArray class with Ray so that it knows how to serialize it.
 ray.register_class(DistArray)
 
+
 @ray.remote
 def assemble(a):
   return a.assemble()
 
-# TODO(rkn): what should we call this method
+
+# TODO(rkn): What should we call this method?
 @ray.remote
 def numpy_to_dist(a):
   result = DistArray(a.shape)
   for index in np.ndindex(*result.num_blocks):
     lower = DistArray.compute_block_lower(index, a.shape)
     upper = DistArray.compute_block_upper(index, a.shape)
-    result.objectids[index] = ray.put(a[[slice(l, u) for (l, u) in zip(lower, upper)]])
+    result.objectids[index] = ray.put(a[[slice(l, u) for (l, u)
+                                        in zip(lower, upper)]])
   return result
 
+
 @ray.remote
 def zeros(shape, dtype_name="float"):
   result = DistArray(shape)
   for index in np.ndindex(*result.num_blocks):
-    result.objectids[index] = ra.zeros.remote(DistArray.compute_block_shape(index, shape), dtype_name=dtype_name)
+    result.objectids[index] = ra.zeros.remote(
+        DistArray.compute_block_shape(index, shape), dtype_name=dtype_name)
   return result
 
+
 @ray.remote
 def ones(shape, dtype_name="float"):
   result = DistArray(shape)
   for index in np.ndindex(*result.num_blocks):
-    result.objectids[index] = ra.ones.remote(DistArray.compute_block_shape(index, shape), dtype_name=dtype_name)
+    result.objectids[index] = ra.ones.remote(
+        DistArray.compute_block_shape(index, shape), dtype_name=dtype_name)
   return result
 
+
 @ray.remote
 def copy(a):
   result = DistArray(a.shape)
   for index in np.ndindex(*result.num_blocks):
-    result.objectids[index] = a.objectids[index] # We don't need to actually copy the objects because cluster-level objects are assumed to be immutable.
+    # We don't need to actually copy the objects because remote objects are
+    # immutable.
+    result.objectids[index] = a.objectids[index]
   return result
 
+
 @ray.remote
 def eye(dim1, dim2=-1, dtype_name="float"):
   dim2 = dim1 if dim2 == -1 else dim2
@@ -107,15 +129,19 @@ def eye(dim1, dim2=-1, dtype_name="float"):
   for (i, j) in np.ndindex(*result.num_blocks):
     block_shape = DistArray.compute_block_shape([i, j], shape)
     if i == j:
-      result.objectids[i, j] = ra.eye.remote(block_shape[0], block_shape[1], dtype_name=dtype_name)
+      result.objectids[i, j] = ra.eye.remote(block_shape[0], block_shape[1],
+                                             dtype_name=dtype_name)
     else:
-      result.objectids[i, j] = ra.zeros.remote(block_shape, dtype_name=dtype_name)
+      result.objectids[i, j] = ra.zeros.remote(block_shape,
+                                               dtype_name=dtype_name)
   return result
 
+
 @ray.remote
 def triu(a):
   if a.ndim != 2:
-    raise Exception("Input must have 2 dimensions, but a.ndim is " + str(a.ndim))
+    raise Exception("Input must have 2 dimensions, but a.ndim is "
+                    "{}.".format(a.ndim))
   result = DistArray(a.shape)
   for (i, j) in np.ndindex(*result.num_blocks):
     if i < j:
@@ -126,10 +152,12 @@ def triu(a):
       result.objectids[i, j] = ra.zeros_like.remote(a.objectids[i, j])
   return result
 
+
 @ray.remote
 def tril(a):
   if a.ndim != 2:
-    raise Exception("Input must have 2 dimensions, but a.ndim is " + str(a.ndim))
+    raise Exception("Input must have 2 dimensions, but a.ndim is "
+                    "{}.".format(a.ndim))
   result = DistArray(a.shape)
   for (i, j) in np.ndindex(*result.num_blocks):
     if i > j:
@@ -140,25 +168,31 @@ def tril(a):
       result.objectids[i, j] = ra.zeros_like.remote(a.objectids[i, j])
   return result
 
+
 @ray.remote
 def blockwise_dot(*matrices):
   n = len(matrices)
   if n % 2 != 0:
-    raise Exception("blockwise_dot expects an even number of arguments, but len(matrices) is {}.".format(n))
+    raise Exception("blockwise_dot expects an even number of arguments, but "
+                    "len(matrices) is {}.".format(n))
   shape = (matrices[0].shape[0], matrices[n // 2].shape[1])
   result = np.zeros(shape)
   for i in range(n // 2):
     result += np.dot(matrices[i], matrices[n // 2 + i])
   return result
 
+
 @ray.remote
 def dot(a, b):
   if a.ndim != 2:
-    raise Exception("dot expects its arguments to be 2-dimensional, but a.ndim = {}.".format(a.ndim))
+    raise Exception("dot expects its arguments to be 2-dimensional, but "
+                    "a.ndim = {}.".format(a.ndim))
   if b.ndim != 2:
-    raise Exception("dot expects its arguments to be 2-dimensional, but b.ndim = {}.".format(b.ndim))
+    raise Exception("dot expects its arguments to be 2-dimensional, but "
+                    "b.ndim = {}.".format(b.ndim))
   if a.shape[1] != b.shape[0]:
-    raise Exception("dot expects a.shape[1] to equal b.shape[0], but a.shape = {} and b.shape = {}.".format(a.shape, b.shape))
+    raise Exception("dot expects a.shape[1] to equal b.shape[0], but a.shape "
+                    "= {} and b.shape = {}.".format(a.shape, b.shape))
   shape = [a.shape[0], b.shape[1]]
   result = DistArray(shape)
   for (i, j) in np.ndindex(*result.num_blocks):
@@ -166,10 +200,12 @@ def dot(a, b):
     result.objectids[i, j] = blockwise_dot.remote(*args)
   return result
 
+
 @ray.remote
 def subblocks(a, *ranges):
   """
-  This function produces a distributed array from a subset of the blocks in the `a`. The result and `a` will have the same number of dimensions.For example,
+  This function produces a distributed array from a subset of the blocks in the
+  `a`. The result and `a` will have the same number of dimensions.For example,
       subblocks(a, [0, 1], [2, 4])
   will produce a DistArray whose objectids are
       [[a.objectids[0, 2], a.objectids[0, 4]],
@@ -178,50 +214,71 @@ def subblocks(a, *ranges):
   """
   ranges = list(ranges)
   if len(ranges) != a.ndim:
-    raise Exception("sub_blocks expects to receive a number of ranges equal to a.ndim, but it received {} ranges and a.ndim = {}.".format(len(ranges), a.ndim))
+    raise Exception("sub_blocks expects to receive a number of ranges equal "
+                    "to a.ndim, but it received {} ranges and a.ndim = "
+                    "{}.".format(len(ranges), a.ndim))
   for i in range(len(ranges)):
-    if ranges[i] == []: # We allow the user to pass in an empty list to indicate the full range
+    # We allow the user to pass in an empty list to indicate the full range.
+    if ranges[i] == []:
       ranges[i] = range(a.num_blocks[i])
     if not np.alltrue(ranges[i] == np.sort(ranges[i])):
-      raise Exception("Ranges passed to sub_blocks must be sorted, but the {}th range is {}.".format(i, ranges[i]))
+      raise Exception("Ranges passed to sub_blocks must be sorted, but the "
+                      "{}th range is {}.".format(i, ranges[i]))
     if ranges[i][0] < 0:
-      raise Exception("Values in the ranges passed to sub_blocks must be at least 0, but the {}th range is {}.".format(i, ranges[i]))
+      raise Exception("Values in the ranges passed to sub_blocks must be at "
+                      "least 0, but the {}th range is {}.".format(i,
+                                                                  ranges[i]))
     if ranges[i][-1] >= a.num_blocks[i]:
-      raise Exception("Values in the ranges passed to sub_blocks must be less than the relevant number of blocks, but the {}th range is {}, and a.num_blocks = {}.".format(i, ranges[i], a.num_blocks))
+      raise Exception("Values in the ranges passed to sub_blocks must be less "
+                      "than the relevant number of blocks, but the {}th range "
+                      "is {}, and a.num_blocks = {}.".format(i, ranges[i],
+                                                             a.num_blocks))
   last_index = [r[-1] for r in ranges]
   last_block_shape = DistArray.compute_block_shape(last_index, a.shape)
-  shape = [(len(ranges[i]) - 1) * BLOCK_SIZE + last_block_shape[i] for i in range(a.ndim)]
+  shape = [(len(ranges[i]) - 1) * BLOCK_SIZE + last_block_shape[i]
+           for i in range(a.ndim)]
   result = DistArray(shape)
   for index in np.ndindex(*result.num_blocks):
-    result.objectids[index] = a.objectids[tuple([ranges[i][index[i]] for i in range(a.ndim)])]
+    result.objectids[index] = a.objectids[tuple([ranges[i][index[i]]
+                                          for i in range(a.ndim)])]
   return result
 
+
 @ray.remote
 def transpose(a):
   if a.ndim != 2:
-    raise Exception("transpose expects its argument to be 2-dimensional, but a.ndim = {}, a.shape = {}.".format(a.ndim, a.shape))
+    raise Exception("transpose expects its argument to be 2-dimensional, but "
+                    "a.ndim = {}, a.shape = {}.".format(a.ndim, a.shape))
   result = DistArray([a.shape[1], a.shape[0]])
   for i in range(result.num_blocks[0]):
     for j in range(result.num_blocks[1]):
       result.objectids[i, j] = ra.transpose.remote(a.objectids[j, i])
   return result
 
+
 # TODO(rkn): support broadcasting?
 @ray.remote
 def add(x1, x2):
   if x1.shape != x2.shape:
-    raise Exception("add expects arguments `x1` and `x2` to have the same shape, but x1.shape = {}, and x2.shape = {}.".format(x1.shape, x2.shape))
+    raise Exception("add expects arguments `x1` and `x2` to have the same "
+                    "shape, but x1.shape = {}, and x2.shape = {}."
+                    .format(x1.shape, x2.shape))
   result = DistArray(x1.shape)
   for index in np.ndindex(*result.num_blocks):
-    result.objectids[index] = ra.add.remote(x1.objectids[index], x2.objectids[index])
+    result.objectids[index] = ra.add.remote(x1.objectids[index],
+                                            x2.objectids[index])
   return result
 
+
 # TODO(rkn): support broadcasting?
 @ray.remote
 def subtract(x1, x2):
   if x1.shape != x2.shape:
-    raise Exception("subtract expects arguments `x1` and `x2` to have the same shape, but x1.shape = {}, and x2.shape = {}.".format(x1.shape, x2.shape))
+    raise Exception("subtract expects arguments `x1` and `x2` to have the "
+                    "same shape, but x1.shape = {}, and x2.shape = {}."
+                    .format(x1.shape, x2.shape))
   result = DistArray(x1.shape)
   for index in np.ndindex(*result.num_blocks):
-    result.objectids[index] = ra.subtract.remote(x1.objectids[index], x2.objectids[index])
+    result.objectids[index] = ra.subtract.remote(x1.objectids[index],
+                                                 x2.objectids[index])
   return result

python/ray/experimental/array/distributed/linalg.py

@@ -6,30 +6,32 @@ import numpy as np
 import ray.experimental.array.remote as ra
 import ray
 
-from .core import *
+from . import core
 
 __all__ = ["tsqr", "modified_lu", "tsqr_hr", "qr"]
 
+
 @ray.remote(num_return_vals=2)
 def tsqr(a):
-  """
-  arguments:
-    a: a distributed matrix
-  Suppose that
-    a.shape == (M, N)
-    K == min(M, N)
-  return values:
-    q: DistArray, if q_full = ray.get(DistArray, q).assemble(), then
-      q_full.shape == (M, K)
-      np.allclose(np.dot(q_full.T, q_full), np.eye(K)) == True
-    r: np.ndarray, if r_val = ray.get(np.ndarray, r), then
-      r_val.shape == (K, N)
-      np.allclose(r, np.triu(r)) == True
+  """Perform a QR decomposition of a tall-skinny matrix.
+
+  Args:
+    a: A distributed matrix with shape MxN (suppose K = min(M, N)).
+
+  Returns:
+    A tuple of q (a DistArray) and r (a numpy array) satisfying the following.
+      - If q_full = ray.get(DistArray, q).assemble(), then
+        q_full.shape == (M, K).
+      - np.allclose(np.dot(q_full.T, q_full), np.eye(K)) == True.
+      - If r_val = ray.get(np.ndarray, r), then r_val.shape == (K, N).
+      - np.allclose(r, np.triu(r)) == True.
   """
   if len(a.shape) != 2:
-    raise Exception("tsqr requires len(a.shape) == 2, but a.shape is {}".format(a.shape))
+    raise Exception("tsqr requires len(a.shape) == 2, but a.shape is "
+                    "{}".format(a.shape))
   if a.num_blocks[1] != 1:
-    raise Exception("tsqr requires a.num_blocks[1] == 1, but a.num_blocks is {}".format(a.num_blocks))
+    raise Exception("tsqr requires a.num_blocks[1] == 1, but a.num_blocks is "
+                    "{}".format(a.num_blocks))
 
   num_blocks = a.num_blocks[0]
   K = int(np.ceil(np.log2(num_blocks))) + 1
@@ -57,9 +59,9 @@ def tsqr(a):
     q_shape = a.shape
   else:
     q_shape = [a.shape[0], a.shape[0]]
-  q_num_blocks = DistArray.compute_num_blocks(q_shape)
+  q_num_blocks = core.DistArray.compute_num_blocks(q_shape)
   q_objectids = np.empty(q_num_blocks, dtype=object)
-  q_result = DistArray(q_shape, q_objectids)
+  q_result = core.DistArray(q_shape, q_objectids)
 
   # reconstruct output
   for i in range(num_blocks):
@@ -68,28 +70,35 @@ def tsqr(a):
     for j in range(1, K):
       if np.mod(ith_index, 2) == 0:
         lower = [0, 0]
-        upper = [a.shape[1], BLOCK_SIZE]
+        upper = [a.shape[1], core.BLOCK_SIZE]
       else:
         lower = [a.shape[1], 0]
-        upper = [2 * a.shape[1], BLOCK_SIZE]
+        upper = [2 * a.shape[1], core.BLOCK_SIZE]
       ith_index //= 2
-      q_block_current = ra.dot.remote(q_block_current, ra.subarray.remote(q_tree[ith_index, j], lower, upper))
+      q_block_current = ra.dot.remote(q_block_current,
+                                      ra.subarray.remote(q_tree[ith_index, j],
+                                                         lower, upper))
     q_result.objectids[i] = q_block_current
   r = current_rs[0]
   return q_result, ray.get(r)
 
+
 # TODO(rkn): This is unoptimized, we really want a block version of this.
+# This is Algorithm 5 from
+# http://www.eecs.berkeley.edu/Pubs/TechRpts/2013/EECS-2013-175.pdf.
 @ray.remote(num_return_vals=3)
 def modified_lu(q):
-  """
-  Algorithm 5 from http://www.eecs.berkeley.edu/Pubs/TechRpts/2013/EECS-2013-175.pdf
-  takes a matrix q with orthonormal columns, returns l, u, s such that q - s = l * u
-  arguments:
-    q: a two dimensional orthonormal q
-  return values:
-    l: lower triangular
-    u: upper triangular
-    s: a diagonal matrix represented by its diagonal
+  """Perform a modified LU decomposition of a matrix.
+
+  This takes a matrix q with orthonormal columns, returns l, u, s such that
+  q - s = l * u.
+
+  Args:
+    q: A two dimensional orthonormal matrix q.
+
+  Returns:
+    A tuple of a lower triangular matrix l, an upper triangular matrix u, and a
+      a vector representing a diagonal matrix s such that q - s = l * u.
   """
   q = q.assemble()
   m, b = q.shape[0], q.shape[1]
@@ -100,14 +109,19 @@ def modified_lu(q):
   for i in range(b):
     S[i] = -1 * np.sign(q_work[i, i])
     q_work[i, i] -= S[i]
-    q_work[(i + 1):m, i] /= q_work[i, i] # scale ith column of L by diagonal element
-    q_work[(i + 1):m, (i + 1):b] -= np.outer(q_work[(i + 1):m, i], q_work[i, (i + 1):b]) # perform Schur complement update
+    # Scale ith column of L by diagonal element.
+    q_work[(i + 1):m, i] /= q_work[i, i]
+    # Perform Schur complement update.
+    q_work[(i + 1):m, (i + 1):b] -= np.outer(q_work[(i + 1):m, i],
+                                             q_work[i, (i + 1):b])
 
   L = np.tril(q_work)
   for i in range(b):
     L[i, i] = 1
   U = np.triu(q_work)[:b, :]
-  return ray.get(numpy_to_dist.remote(ray.put(L))), U, S # TODO(rkn): get rid of put
+  # TODO(rkn): Get rid of the put below.
+  return ray.get(core.numpy_to_dist.remote(ray.put(L))), U, S
+
 
 @ray.remote(num_return_vals=2)
 def tsqr_hr_helper1(u, s, y_top_block, b):
@@ -116,45 +130,60 @@ def tsqr_hr_helper1(u, s, y_top_block, b):
   t = -1 * np.dot(u, np.dot(s_full, np.linalg.inv(y_top).T))
   return t, y_top
 
+
 @ray.remote
 def tsqr_hr_helper2(s, r_temp):
   s_full = np.diag(s)
   return np.dot(s_full, r_temp)
 
+
+# This is Algorithm 6 from
+# http://www.eecs.berkeley.edu/Pubs/TechRpts/2013/EECS-2013-175.pdf.
 @ray.remote(num_return_vals=4)
 def tsqr_hr(a):
-  """Algorithm 6 from http://www.eecs.berkeley.edu/Pubs/TechRpts/2013/EECS-2013-175.pdf"""
   q, r_temp = tsqr.remote(a)
   y, u, s = modified_lu.remote(q)
   y_blocked = ray.get(y)
-  t, y_top = tsqr_hr_helper1.remote(u, s, y_blocked.objectids[0, 0], a.shape[1])
+  t, y_top = tsqr_hr_helper1.remote(u, s, y_blocked.objectids[0, 0],
+                                    a.shape[1])
   r = tsqr_hr_helper2.remote(s, r_temp)
   return ray.get(y), ray.get(t), ray.get(y_top), ray.get(r)
 
+
 @ray.remote
 def qr_helper1(a_rc, y_ri, t, W_c):
   return a_rc - np.dot(y_ri, np.dot(t.T, W_c))
 
+
 @ray.remote
 def qr_helper2(y_ri, a_rc):
   return np.dot(y_ri.T, a_rc)
 
+
+# This is Algorithm 7 from
+# http://www.eecs.berkeley.edu/Pubs/TechRpts/2013/EECS-2013-175.pdf.
 @ray.remote(num_return_vals=2)
 def qr(a):
-  """Algorithm 7 from http://www.eecs.berkeley.edu/Pubs/TechRpts/2013/EECS-2013-175.pdf"""
+
   m, n = a.shape[0], a.shape[1]
   k = min(m, n)
 
   # we will store our scratch work in a_work
-  a_work = DistArray(a.shape, np.copy(a.objectids))
+  a_work = core.DistArray(a.shape, np.copy(a.objectids))
 
   result_dtype = np.linalg.qr(ray.get(a.objectids[0, 0]))[0].dtype.name
-  r_res = ray.get(zeros.remote([k, n], result_dtype)) # TODO(rkn): It would be preferable not to get this right after creating it.
-  y_res = ray.get(zeros.remote([m, k], result_dtype)) # TODO(rkn): It would be preferable not to get this right after creating it.
+  # TODO(rkn): It would be preferable not to get this right after creating it.
+  r_res = ray.get(core.zeros.remote([k, n], result_dtype))
+  # TODO(rkn): It would be preferable not to get this right after creating it.
+  y_res = ray.get(core.zeros.remote([m, k], result_dtype))
   Ts = []
 
-  for i in range(min(a.num_blocks[0], a.num_blocks[1])): # this differs from the paper, which says "for i in range(a.num_blocks[1])", but that doesn't seem to make any sense when a.num_blocks[1] > a.num_blocks[0]
-    sub_dist_array = subblocks.remote(a_work, list(range(i, a_work.num_blocks[0])), [i])
+  # The for loop differs from the paper, which says
+  # "for i in range(a.num_blocks[1])", but that doesn't seem to make any sense
+  # when a.num_blocks[1] > a.num_blocks[0].
+  for i in range(min(a.num_blocks[0], a.num_blocks[1])):
+    sub_dist_array = core.subblocks.remote(
+        a_work, list(range(i, a_work.num_blocks[0])), [i])
     y, t, _, R = tsqr_hr.remote(sub_dist_array)
     y_val = ray.get(y)
 
@@ -167,7 +196,7 @@ def qr(a):
       r_res.objectids[i, i] = ra.dot.remote(eye_temp, R)
     else:
       r_res.objectids[i, i] = R
-    Ts.append(numpy_to_dist.remote(t))
+    Ts.append(core.numpy_to_dist.remote(t))
 
     for c in range(i + 1, a.num_blocks[1]):
       W_rcs = []
@@ -182,9 +211,14 @@ def qr(a):
       r_res.objectids[i, c] = a_work.objectids[i, c]
 
   # construct q_res from Ys and Ts
-  q = eye.remote(m, k, dtype_name=result_dtype)
+  q = core.eye.remote(m, k, dtype_name=result_dtype)
   for i in range(len(Ts))[::-1]:
-    y_col_block = subblocks.remote(y_res, [], [i])
-    q = subtract.remote(q, dot.remote(y_col_block, dot.remote(Ts[i], dot.remote(transpose.remote(y_col_block), q))))
+    y_col_block = core.subblocks.remote(y_res, [], [i])
+    q = core.subtract.remote(
+        q, core.dot.remote(
+            y_col_block,
+            core.dot.remote(Ts[i],
+                            core.dot.remote(core.transpose.remote(y_col_block),
+                                            q))))
 
   return ray.get(q), r_res

python/ray/experimental/array/distributed/random.py

@@ -6,13 +6,15 @@ import numpy as np
 import ray.experimental.array.remote as ra
 import ray
 
-from .core import *
+from .core import DistArray
+
 
 @ray.remote
 def normal(shape):
   num_blocks = DistArray.compute_num_blocks(shape)
   objectids = np.empty(num_blocks, dtype=object)
   for index in np.ndindex(*num_blocks):
-    objectids[index] = ra.random.normal.remote(DistArray.compute_block_shape(index, shape))
+    objectids[index] = ra.random.normal.remote(
+        DistArray.compute_block_shape(index, shape))
   result = DistArray(shape, objectids)
   return result

python/ray/experimental/array/remote/__init__.py

@@ -4,4 +4,10 @@ from __future__ import print_function
 
 from . import random
 from . import linalg
-from .core import *
+from .core import (zeros, zeros_like, ones, eye, dot, vstack, hstack, subarray,
+                   copy, tril, triu, diag, transpose, add, subtract, sum,
+                   shape, sum_list)
+
+__all__ = ["random", "linalg", "zeros", "zeros_like", "ones", "eye", "dot",
+           "vstack", "hstack", "subarray", "copy", "tril", "triu", "diag",
+           "transpose", "add", "subtract", "sum", "shape", "sum_list"]

python/ray/experimental/array/remote/core.py

@@ -5,82 +5,97 @@ from __future__ import print_function
 import numpy as np
 import ray
 
-__all__ = ["zeros", "zeros_like", "ones", "eye", "dot", "vstack", "hstack", "subarray", "copy", "tril", "triu", "diag", "transpose", "add", "subtract", "sum", "shape", "sum_list"]
 
 @ray.remote
 def zeros(shape, dtype_name="float", order="C"):
   return np.zeros(shape, dtype=np.dtype(dtype_name), order=order)
 
+
 @ray.remote
 def zeros_like(a, dtype_name="None", order="K", subok=True):
   dtype_val = None if dtype_name == "None" else np.dtype(dtype_name)
   return np.zeros_like(a, dtype=dtype_val, order=order, subok=subok)
 
+
 @ray.remote
 def ones(shape, dtype_name="float", order="C"):
   return np.ones(shape, dtype=np.dtype(dtype_name), order=order)
 
+
 @ray.remote
 def eye(N, M=-1, k=0, dtype_name="float"):
   M = N if M == -1 else M
   return np.eye(N, M=M, k=k, dtype=np.dtype(dtype_name))
 
+
 @ray.remote
 def dot(a, b):
   return np.dot(a, b)
 
+
 @ray.remote
 def vstack(*xs):
   return np.vstack(xs)
 
+
 @ray.remote
 def hstack(*xs):
   return np.hstack(xs)
 
-# TODO(rkn): instead of this, consider implementing slicing
+
+# TODO(rkn): Instead of this, consider implementing slicing.
+# TODO(rkn): Be consistent about using "index" versus "indices".
 @ray.remote
-def subarray(a, lower_indices, upper_indices): # TODO(rkn): be consistent about using "index" versus "indices"
+def subarray(a, lower_indices, upper_indices):
   return a[[slice(l, u) for (l, u) in zip(lower_indices, upper_indices)]]
 
+
 @ray.remote
 def copy(a, order="K"):
   return np.copy(a, order=order)
 
+
 @ray.remote
 def tril(m, k=0):
   return np.tril(m, k=k)
 
+
 @ray.remote
 def triu(m, k=0):
   return np.triu(m, k=k)
 
+
 @ray.remote
 def diag(v, k=0):
   return np.diag(v, k=k)
 
+
 @ray.remote
 def transpose(a, axes=[]):
   axes = None if axes == [] else axes
   return np.transpose(a, axes=axes)
 
+
 @ray.remote
 def add(x1, x2):
   return np.add(x1, x2)
 
+
 @ray.remote
 def subtract(x1, x2):
   return np.subtract(x1, x2)
 
+
 @ray.remote
 def sum(x, axis=-1):
   return np.sum(x, axis=axis if axis != -1 else None)
 
+
 @ray.remote
 def shape(a):
   return np.shape(a)
 
-# We use Any to allow different numerical types as well as numpy arrays.
-# TODO(rkn):this isn't in the numpy API, so be careful about exposing this.
+
 @ray.remote
 def sum_list(*xs):
   return np.sum(xs, axis=0)

python/ray/experimental/array/remote/linalg.py

@@ -8,84 +8,104 @@ import ray
 __all__ = ["matrix_power", "solve", "tensorsolve", "tensorinv", "inv",
            "cholesky", "eigvals", "eigvalsh", "pinv", "slogdet", "det",
            "svd", "eig", "eigh", "lstsq", "norm", "qr", "cond", "matrix_rank",
-           "LinAlgError", "multi_dot"]
+           "multi_dot"]
+
 
 @ray.remote
 def matrix_power(M, n):
   return np.linalg.matrix_power(M, n)
 
+
 @ray.remote
 def solve(a, b):
   return np.linalg.solve(a, b)
 
+
 @ray.remote(num_return_vals=2)
 def tensorsolve(a):
   raise NotImplementedError
 
+
 @ray.remote(num_return_vals=2)
 def tensorinv(a):
   raise NotImplementedError
 
+
 @ray.remote
 def inv(a):
   return np.linalg.inv(a)
 
+
 @ray.remote
 def cholesky(a):
   return np.linalg.cholesky(a)
 
+
 @ray.remote
 def eigvals(a):
   return np.linalg.eigvals(a)
 
+
 @ray.remote
 def eigvalsh(a):
   raise NotImplementedError
 
+
 @ray.remote
 def pinv(a):
   return np.linalg.pinv(a)
 
+
 @ray.remote
 def slogdet(a):
   raise NotImplementedError
 
+
 @ray.remote
 def det(a):
   return np.linalg.det(a)
 
+
 @ray.remote(num_return_vals=3)
 def svd(a):
   return np.linalg.svd(a)
 
+
 @ray.remote(num_return_vals=2)
 def eig(a):
   return np.linalg.eig(a)
 
+
 @ray.remote(num_return_vals=2)
 def eigh(a):
   return np.linalg.eigh(a)
 
+
 @ray.remote(num_return_vals=4)
 def lstsq(a, b):
   return np.linalg.lstsq(a)
 
+
 @ray.remote
 def norm(x):
   return np.linalg.norm(x)
 
+
 @ray.remote(num_return_vals=2)
 def qr(a):
   return np.linalg.qr(a)
 
+
 @ray.remote
 def cond(x):
   return np.linalg.cond(x)
 
+
 @ray.remote
 def matrix_rank(M):
   return np.linalg.matrix_rank(M)
 
+
 @ray.remote
 def multi_dot(*a):
   raise NotImplementedError

python/ray/experimental/array/remote/random.py

@@ -5,6 +5,7 @@ from __future__ import print_function
 import numpy as np
 import ray
 
+
 @ray.remote
 def normal(shape):
   return np.random.normal(size=shape)