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)
