mirror of
https://github.com/wassname/ray.git
synced 2026-07-12 06:19:23 +08:00
rewrite linear algebra libraries to use keyword arguments (#78)
This commit is contained in:
Robert Nishihara
committed by
Philipp Moritz
Philipp Moritz
parent
0aa1fa097e
commit
b58eaf84ee
Vendored
+9
-20
@@ -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])
|
||||
|
||||
Vendored
+2
-2
@@ -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))))
|
||||
|
||||
@@ -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):
|
||||
|
||||
+6
-6
@@ -25,12 +25,12 @@ class ArraysSingleTest(unittest.TestCase):
|
||||
services.start_singlenode_cluster(return_drivers=False, num_workers_per_objstore=1, worker_path=test_path)
|
||||
|
||||
# test eye
|
||||
ref = single.eye(3, "float")
|
||||
ref = single.eye(3)
|
||||
val = orchpy.pull(ref)
|
||||
self.assertTrue(np.alltrue(val == np.eye(3)))
|
||||
|
||||
# test zeros
|
||||
ref = single.zeros([3, 4, 5], "float")
|
||||
ref = single.zeros([3, 4, 5])
|
||||
val = orchpy.pull(ref)
|
||||
self.assertTrue(np.alltrue(val == np.zeros([3, 4, 5])))
|
||||
|
||||
@@ -70,8 +70,8 @@ class ArraysDistTest(unittest.TestCase):
|
||||
test_path = os.path.join(test_dir, "testrecv.py")
|
||||
services.start_singlenode_cluster(return_drivers=False, num_workers_per_objstore=1, worker_path=test_path)
|
||||
|
||||
a = single.ones([dist.BLOCK_SIZE, dist.BLOCK_SIZE], "float")
|
||||
b = single.zeros([dist.BLOCK_SIZE, dist.BLOCK_SIZE], "float")
|
||||
a = single.ones([dist.BLOCK_SIZE, dist.BLOCK_SIZE])
|
||||
b = single.zeros([dist.BLOCK_SIZE, dist.BLOCK_SIZE])
|
||||
x = dist.DistArray()
|
||||
x.construct([2 * dist.BLOCK_SIZE, dist.BLOCK_SIZE], np.array([[a], [b]]))
|
||||
self.assertTrue(np.alltrue(x.assemble() == np.vstack([np.ones([dist.BLOCK_SIZE, dist.BLOCK_SIZE]), np.zeros([dist.BLOCK_SIZE, dist.BLOCK_SIZE])])))
|
||||
@@ -87,7 +87,7 @@ class ArraysDistTest(unittest.TestCase):
|
||||
y = dist.assemble(x)
|
||||
self.assertTrue(np.alltrue(orchpy.pull(y) == np.zeros([9, 25, 51])))
|
||||
|
||||
x = dist.ones([11, 25, 49], "float")
|
||||
x = dist.ones([11, 25, 49], dtype_name="float")
|
||||
y = dist.assemble(x)
|
||||
self.assertTrue(np.alltrue(orchpy.pull(y) == np.ones([11, 25, 49])))
|
||||
|
||||
@@ -97,7 +97,7 @@ class ArraysDistTest(unittest.TestCase):
|
||||
w = dist.assemble(y)
|
||||
self.assertTrue(np.alltrue(orchpy.pull(z) == orchpy.pull(w)))
|
||||
|
||||
x = dist.eye(25, "float")
|
||||
x = dist.eye(25, dtype_name="float")
|
||||
y = dist.assemble(x)
|
||||
self.assertTrue(np.alltrue(orchpy.pull(y) == np.eye(25)))
|
||||
|
||||
|
||||
Reference in New Issue
Block a user