def cholesky(self, X: BlockArray): # TODO (hme): Implement scalable version. # Note: # A = Q, R # A.T @ A = R.T @ R # A.T @ A = L @ L.T # => R == L.T block_shape = X.block_shape assert len(X.shape) == 2 assert X.shape[0] == X.shape[1] single_block = X.shape[0] == X.block_shape[0] and X.shape[ 1] == X.block_shape[1] if single_block: result = X.copy() else: result = X.reshape(block_shape=X.shape) result.blocks[0, 0].oid = self._system.cholesky(result.blocks[0, 0].oid, syskwargs={ "grid_entry": (0, 0), "grid_shape": (1, 1) }) if not single_block: result = result.reshape(block_shape=block_shape) return result
def indirect_tsr(self, X: BlockArray, reshape_output=True): assert len(X.shape) == 2 # TODO (hme): This assertion is temporary and ensures returned # shape of qr of block is correct. assert X.block_shape[0] >= X.shape[1] # Compute R for each block. grid = X.grid grid_shape = grid.grid_shape shape = X.shape block_shape = X.block_shape R_oids = [] # Assume no blocking along second dim. for i in range(grid_shape[0]): # Select a row according to block_shape. row = [] for j in range(grid_shape[1]): row.append(X.blocks[i, j].oid) R_oids.append( self._system.qr(*row, mode="r", axis=1, syskwargs={ "grid_entry": (i, 0), "grid_shape": (grid_shape[0], 1), "options": { "num_return_vals": 1 } })) # Construct R by summing over R blocks. # TODO (hme): Communication may be inefficient due to redundancy of data. R_shape = (shape[1], shape[1]) R_block_shape = (block_shape[1], block_shape[1]) tsR = BlockArray( ArrayGrid(shape=R_shape, block_shape=R_shape, dtype=X.dtype.__name__), self._system) tsR.blocks[0, 0].oid = self._system.qr(*R_oids, mode="r", axis=0, syskwargs={ "grid_entry": (0, 0), "grid_shape": (1, 1), "options": { "num_return_vals": 1 } }) # If blocking is "tall-skinny," then we're done. if R_shape != R_block_shape: if reshape_output: R = tsR.reshape(shape=R_shape, block_shape=R_block_shape) else: R = tsR else: R = tsR return R
def vec_from_oids(self, oids, shape, block_shape, dtype): arr = BlockArray( ArrayGrid(shape=shape, block_shape=shape, dtype=dtype.__name__), self.cm) # Make sure resulting grid shape is a vector (1 dimensional). assert np.sum(arr.grid.grid_shape) == (max(arr.grid.grid_shape) + len(arr.grid.grid_shape) - 1) for i, grid_entry in enumerate(arr.grid.get_entry_iterator()): arr.blocks[grid_entry].oid = oids[i] if block_shape != shape: return arr.reshape(block_shape=block_shape) return arr
def permutation(self, size, block_size): shape = (size,) block_shape = (block_size,) grid: ArrayGrid = ArrayGrid(shape=shape, block_shape=shape, dtype=np.int64.__name__) ba = BlockArray(grid, self._system) for grid_entry in ba.grid.get_entry_iterator(): rng_params = list(self._rng.new_block_rng_params()) block: Block = ba.blocks[grid_entry] block.oid = self._system.permutation(rng_params, size, syskwargs={ "grid_entry": grid_entry, "grid_shape": grid.grid_shape }) return ba.reshape(block_shape=block_shape)
def inv(app: ArrayApplication, X: BlockArray): # TODO (hme): Implement scalable version. block_shape = X.block_shape assert len(X.shape) == 2 assert X.shape[0] == X.shape[1] single_block = X.shape[0] == X.block_shape[0] and X.shape[1] == X.block_shape[1] if single_block: result = X.copy() else: result = X.reshape(block_shape=X.shape) result.blocks[0, 0].oid = app.cm.inv(result.blocks[0, 0].oid, syskwargs={ "grid_entry": (0, 0), "grid_shape": (1, 1) }) if not single_block: result = result.reshape(block_shape=block_shape) return result
def _inv(self, remote_func, kwargs, X: BlockArray): # TODO (hme): Implement scalable version. block_shape = X.block_shape assert len(X.shape) == 2 assert X.shape[0] == X.shape[1] single_block = X.shape[0] == X.block_shape[0] and X.shape[ 1] == X.block_shape[1] if single_block: result = X.copy() else: result = X.reshape(block_shape=X.shape) result.blocks[0, 0].oid = remote_func(result.blocks[0, 0].oid, **kwargs, syskwargs={ "grid_entry": (0, 0), "grid_shape": (1, 1) }) if not single_block: result = result.reshape(block_shape=block_shape) return result
def inv_sym_psd(self, X: BlockArray): # Assumes X is symmetric PSD. # TODO (hme): Implement scalable version. assert len(X.shape) == 2 assert X.shape[0] == X.shape[1] single_block = X.shape[0] == X.block_shape[0] and X.shape[ 1] == X.block_shape[1] if single_block: result = X.copy() else: result = X.reshape(block_shape=X.shape) result.blocks[0, 0].oid = self._system.inv_sym_psd(result.blocks[0, 0].oid, syskwargs={ "grid_entry": (0, 0), "grid_shape": (1, 1) }) if not single_block: result = result.reshape(block_shape=X.block_shape) return result
def reshape(a: BlockArray, shape): block_shape = _instance().compute_block_shape(shape, a.dtype) return a.reshape(shape, block_shape=block_shape)
def outer(a: BlockArray, b: BlockArray): assert len(a.shape) == len( b.shape) == 1, "Only single-axis inputs supported." return a.reshape((a.shape[0], 1)) @ b.reshape((1, b.shape[0]))
def reshape(x: BlockArray, shape): block_shape = _instance().compute_block_shape(shape, x.dtype) return x.reshape(shape, block_shape=block_shape)