diff --git a/lib/orchpy/arrays/dist/core.py b/lib/orchpy/arrays/dist/core.py index 0e22d37e3..37b54d41f 100644 --- a/lib/orchpy/arrays/dist/core.py +++ b/lib/orchpy/arrays/dist/core.py @@ -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]) diff --git a/lib/orchpy/arrays/dist/linalg.py b/lib/orchpy/arrays/dist/linalg.py index 1d9eba246..9e58d3c62 100644 --- a/lib/orchpy/arrays/dist/linalg.py +++ b/lib/orchpy/arrays/dist/linalg.py @@ -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)))) diff --git a/lib/orchpy/arrays/single/core.py b/lib/orchpy/arrays/single/core.py index 05ca1a545..227e5fcfd 100644 --- a/lib/orchpy/arrays/single/core.py +++ b/lib/orchpy/arrays/single/core.py @@ -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): diff --git a/test/arrays_test.py b/test/arrays_test.py index 18e53486b..7b9a9d720 100644 --- a/test/arrays_test.py +++ b/test/arrays_test.py @@ -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)))