wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-08-19 12:30:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

rewrite linear algebra libraries to use keyword arguments (#78)

Browse Source
This commit is contained in:
Robert Nishihara
2016-06-03 17:09:52 -07:00
committed by Philipp Moritz
parent 0aa1fa097e
commit b58eaf84ee
4 changed files with 48 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
lib/orchpy/arrays/dist/core.py
Vendored
+9 -20
View File
@@ -4,7 +4,7 @@ import arrays.single as single
import orchpy as op
__all__ = ["BLOCK_SIZE", "DistArray", "assemble", "zeros", "ones", "copy",
"eye", "triu", "tril", "blockwise_dot", "dot", "transpose", "add", "subtract", "eye2", "numpy_to_dist", "subblocks"]
"eye", "triu", "tril", "blockwise_dot", "dot", "transpose", "add", "subtract", "numpy_to_dist", "subblocks"]
BLOCK_SIZE = 10
@@ -84,17 +84,17 @@ def numpy_to_dist(a):
return result
@op.distributed([List[int], str], [DistArray])
def zeros(shape, dtype_name):
def zeros(shape, dtype_name="float"):
result = DistArray(shape)
for index in np.ndindex(*result.num_blocks):
result.objrefs[index] = single.zeros(DistArray.compute_block_shape(index, shape), dtype_name)
result.objrefs[index] = single.zeros(DistArray.compute_block_shape(index, shape), dtype_name=dtype_name)
return result
@op.distributed([List[int], str], [DistArray])
def ones(shape, dtype_name):
def ones(shape, dtype_name="float"):
result = DistArray(shape)
for index in np.ndindex(*result.num_blocks):
result.objrefs[index] = single.ones(DistArray.compute_block_shape(index, shape), dtype_name)
result.objrefs[index] = single.ones(DistArray.compute_block_shape(index, shape), dtype_name=dtype_name)
return result
@op.distributed([DistArray], [DistArray])
@@ -104,28 +104,17 @@ def copy(a):
result.objrefs[index] = a.objrefs[index] # We don't need to actually copy the objects because cluster-level objects are assumed to be immutable.
return result
@op.distributed([int, str], [DistArray])
def eye(dim, dtype_name):
shape = [dim, dim]
result = DistArray(shape)
for (i, j) in np.ndindex(*result.num_blocks):
if i == j:
result.objrefs[i, j] = single.eye(DistArray.compute_block_shape([i, j], shape)[0], dtype_name)
else:
result.objrefs[i, j] = single.zeros(DistArray.compute_block_shape([i, j], shape), dtype_name)
return result
# TODO(rkn): Support optional arguments so that we can make this part of eye.
@op.distributed([int, int, str], [DistArray])
def eye2(dim1, dim2, dtype_name):
def eye(dim1, dim2=-1, dtype_name="float"):
dim2 = dim1 if dim2 == -1 else dim2
shape = [dim1, dim2]
result = DistArray(shape)
for (i, j) in np.ndindex(*result.num_blocks):
block_shape = DistArray.compute_block_shape([i, j], shape)
if i == j:
result.objrefs[i, j] = single.eye2(block_shape[0], block_shape[1], dtype_name)
result.objrefs[i, j] = single.eye(block_shape[0], block_shape[1], dtype_name=dtype_name)
else:
result.objrefs[i, j] = single.zeros(block_shape, dtype_name)
result.objrefs[i, j] = single.zeros(block_shape, dtype_name=dtype_name)
return result
@op.distributed([DistArray], [DistArray])
lib/orchpy/arrays/dist/linalg.py
Vendored
+2 -2
View File
@@ -165,7 +165,7 @@ def qr(a):
if a.shape[0] > a.shape[1]:
# in this case, R needs to be square
R_shape = op.pull(single.shape(R))
eye_temp = single.eye2(R_shape[1], R_shape[0], result_dtype)
eye_temp = single.eye(R_shape[1], R_shape[0], dtype_name=result_dtype)
r_res.objrefs[i, i] = single.dot(eye_temp, R)
else:
r_res.objrefs[i, i] = R
@@ -184,7 +184,7 @@ def qr(a):
r_res.objrefs[i, c] = a_work.objrefs[i, c]
# construct q_res from Ys and Ts
q = eye2(m, k, result_dtype)
q = eye(m, k, dtype_name=result_dtype)
for i in range(len(Ts))[::-1]:
y_col_block = subblocks(y_res, [], [i])
q = subtract(q, dot(y_col_block, dot(Ts[i], dot(transpose(y_col_block), q))))
lib/orchpy/arrays/single/core.py
+31 -33
View File
@@ -2,28 +2,25 @@ from typing import List
import numpy as np
import orchpy as op
__all__ = ["zeros", "zeros_like", "ones", "eye", "dot", "vstack", "hstack", "subarray", "copy", "tril", "triu", "diag", "transpose", "add", "subtract", "eye2", "sum", "shape"]
__all__ = ["zeros", "zeros_like", "ones", "eye", "dot", "vstack", "hstack", "subarray", "copy", "tril", "triu", "diag", "transpose", "add", "subtract", "sum", "shape"]
@op.distributed([List[int], str], [np.ndarray])
def zeros(shape, dtype_name):
return np.zeros(shape, dtype=np.dtype(dtype_name))
@op.distributed([List[int], str, str], [np.ndarray])
def zeros(shape, dtype_name="float", order="C"):
return np.zeros(shape, dtype=np.dtype(dtype_name), order=order)
@op.distributed([np.ndarray], [np.ndarray])
def zeros_like(x):
return np.zeros_like(x)
@op.distributed([np.ndarray, str, str, bool], [np.ndarray])
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)
@op.distributed([List[int], str], [np.ndarray])
def ones(shape, dtype_name):
return np.ones(shape, dtype=np.dtype(dtype_name))
@op.distributed([List[int], str, str], [np.ndarray])
def ones(shape, dtype_name="float", order="C"):
return np.ones(shape, dtype=np.dtype(dtype_name), order=order)
@op.distributed([int, str], [np.ndarray])
def eye(dim, dtype_name):
return np.eye(dim, dtype=np.dtype(dtype_name))
# TODO(rkn): This should be part of eye
@op.distributed([int, int, str], [np.ndarray])
def eye2(dim1, dim2, dtype_name):
return np.eye(dim1, dim2, dtype=np.dtype(dtype_name))
@op.distributed([int, int, int, str], [np.ndarray])
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))
@op.distributed([np.ndarray, np.ndarray], [np.ndarray])
def dot(a, b):
@@ -45,25 +42,26 @@ def hstack(*xs):
def subarray(a, lower_indices, upper_indices): # TODO(rkn): be consistent about using "index" versus "indices"
return a[[slice(l, u) for (l, u) in zip(lower_indices, upper_indices)]]
@op.distributed([np.ndarray], [np.ndarray])
def copy(a):
return np.copy(a)
@op.distributed([np.ndarray, str], [np.ndarray])
def copy(a, order="K"):
return np.copy(a, order=order)
@op.distributed([np.ndarray], [np.ndarray])
def tril(a):
return np.tril(a)
@op.distributed([np.ndarray, int], [np.ndarray])
def tril(m, k=0):
return np.tril(m, k=k)
@op.distributed([np.ndarray], [np.ndarray])
def triu(a):
return np.triu(a)
@op.distributed([np.ndarray, int], [np.ndarray])
def triu(m, k=0):
return np.triu(m, k=k)
@op.distributed([np.ndarray], [np.ndarray])
def diag(a):
return np.diag(a)
@op.distributed([np.ndarray, int], [np.ndarray])
def diag(v, k=0):
return np.diag(v, k=k)
@op.distributed([np.ndarray], [np.ndarray])
def transpose(a):
return np.transpose(a)
@op.distributed([np.ndarray, List[int]], [np.ndarray])
def transpose(a, axes=[]):
axes = None if axes == [] else axes
return np.transpose(a, axes=axes)
@op.distributed([np.ndarray, np.ndarray], [np.ndarray])
def add(x1, x2):