def __init__(self):
matrix_pairs = {}
rng = np.random.RandomState(SEED)
for outer_dim in SIZES:
for inner_dim in [max(1, outer_dim // 2), max(1, outer_dim - 1)]:
factor_matrix = rng.standard_normal((outer_dim, inner_dim))
inner_pos_def_matrix = rng.standard_normal(
(inner_dim, inner_dim))
inner_pos_def_matrix = (
inner_pos_def_matrix @ inner_pos_def_matrix.T)
pos_def_matrix = rng.standard_normal((outer_dim, outer_dim))
pos_def_matrix = pos_def_matrix @ pos_def_matrix.T
matrix_pairs[(inner_dim, outer_dim)] = (
matrices.PositiveDefiniteLowRankUpdateMatrix(
matrices.DenseRectangularMatrix(factor_matrix),
matrices.DensePositiveDefiniteMatrix(pos_def_matrix),
matrices.DensePositiveDefiniteMatrix(
inner_pos_def_matrix)), pos_def_matrix +
factor_matrix @ (inner_pos_def_matrix @ factor_matrix.T))
if AUTOGRAD_AVAILABLE:
def param_func(param, matrix):
return (matrix.pos_def_matrix.array +
param @ matrix.inner_pos_def_matrix @ param.T)
def get_param(matrix):
return matrix.factor_matrix.array
else:
param_func, get_param = None, None
super().__init__(matrix_pairs, get_param, param_func, rng)
def matrix_pair(self, rng, size, size_inner):
factor_matrix = rng.standard_normal((size, size_inner))
inner_pos_def_matrix = rng.standard_normal((size_inner, size_inner))
inner_pos_def_matrix = inner_pos_def_matrix @ inner_pos_def_matrix.T
pos_def_matrix = rng.standard_normal((size, size))
pos_def_matrix = pos_def_matrix @ pos_def_matrix.T
return (
matrices.PositiveDefiniteLowRankUpdateMatrix(
matrices.DenseRectangularMatrix(factor_matrix),
matrices.DensePositiveDefiniteMatrix(pos_def_matrix),
matrices.DensePositiveDefiniteMatrix(inner_pos_def_matrix),
),
pos_def_matrix + factor_matrix @ (inner_pos_def_matrix @ factor_matrix.T),
)
def _generate_metrics(rng, size, eigval_scale=0.1):
eigval, eigvec = _generate_rand_eigval_eigvec(rng, size, eigval_scale)
return [
matrices.IdentityMatrix(),
matrices.IdentityMatrix(size),
matrices.PositiveDiagonalMatrix(eigval),
matrices.DensePositiveDefiniteMatrix((eigvec * eigval) @ eigvec.T),
matrices.EigendecomposedPositiveDefiniteMatrix(eigvec, eigval)
]
def matrix_pair(self, rng, size, n_block):
arrays = [rng.standard_normal((size, size)) for _ in range(n_block)]
arrays = [arr @ arr.T for arr in arrays]
return (
matrices.PositiveDefiniteBlockDiagonalMatrix(
matrices.DensePositiveDefiniteMatrix(arr) for arr in arrays
),
sla.block_diag(*arrays),
)
def metric_list(rng, size):
eigval = np.exp(0.1 * rng.standard_normal(size))
eigvec = np.linalg.qr(rng.standard_normal((size, size)))[0]
return [
matrices.IdentityMatrix(),
matrices.IdentityMatrix(size),
matrices.PositiveDiagonalMatrix(eigval),
matrices.DensePositiveDefiniteMatrix((eigvec * eigval) @ eigvec.T),
matrices.EigendecomposedPositiveDefiniteMatrix(eigvec, eigval),
]
def __init__(self):
matrix_pairs = {}
rng = np.random.RandomState(SEED)
for s in SIZES:
for n_block in [1, 2, 5]:
arrays = [rng.standard_normal((s, s)) for _ in range(n_block)]
arrays = [arr @ arr.T for arr in arrays]
matrix_pairs[(s, n_block)] = (
matrices.PositiveDefiniteBlockDiagonalMatrix(
matrices.DensePositiveDefiniteMatrix(arr)
for arr in arrays), sla.block_diag(*arrays))
super().__init__(matrix_pairs, rng)
def __init__(self):
super().__init__()
if AUTOGRAD_AVAILABLE:
grad_log_abs_det_sqrt_func = grad(
lambda a: 0.5 * anp.linalg.slogdet(a)[1])
grad_quadratic_form_inv_func = grad(
lambda a, v: v @ anp.linalg.solve(a, v))
for sz in SIZES:
sqrt = self.rng.standard_normal((sz, sz))
array = sqrt @ sqrt.T
self.matrices[sz] = matrices.DensePositiveDefiniteMatrix(array)
self.np_matrices[sz] = array
if AUTOGRAD_AVAILABLE:
self.grad_log_abs_det_sqrts[sz] = (
grad_log_abs_det_sqrt_func(array))
self.grad_quadratic_form_invs[sz] = partial(
grad_quadratic_form_inv_func, array)
def __init__(self):
matrix_pairs = {}
rng = np.random.RandomState(SEED)
for s in SIZES:
for n_block in [1, 2, 5]:
arrays = [rng.standard_normal((s, s)) for _ in range(n_block)]
arrays = [arr @ arr.T for arr in arrays]
matrix_pairs[(s, n_block)] = (
matrices.PositiveDefiniteBlockDiagonalMatrix(
matrices.DensePositiveDefiniteMatrix(arr)
for arr in arrays), sla.block_diag(*arrays))
if AUTOGRAD_AVAILABLE:
@primitive
def block_diag(blocks):
return sla.block_diag(*blocks)
def vjp_block_diag(ans, blocks):
blocks = tuple(blocks)
def vjp(g):
i, j = 0, 0
vjp_blocks = []
for block in blocks:
j += block.shape[0]
vjp_blocks.append(g[i:j, i:j])
i = j
return tuple(vjp_blocks)
return vjp
defvjp(block_diag, vjp_block_diag)
def get_param(matrix):
return tuple(block.array for block in matrix._blocks)
def param_func(param, matrix):
return block_diag(param)
else:
param_func, get_param = None, None
super().__init__(matrix_pairs, get_param, param_func, rng)
def __init__(self):
super().__init__(
lambda array, is_posdef: matrices.DensePositiveDefiniteMatrix(
array), (+1, ))
def matrix_class(array, is_posdef):
return matrices.DensePositiveDefiniteMatrix(array)