def test_batch_cg_with_tridiag(self):
batch = 5
size = 10
matrix = torch.randn(batch, size, size, dtype=torch.float64)
matrix = matrix.matmul(matrix.transpose(-1, -2))
matrix.div_(matrix.norm())
matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))
rhs = torch.randn(batch, size, 50, dtype=torch.float64)
solves, t_mats = linear_cg(matrix.matmul,
rhs=rhs,
n_tridiag=8,
max_iter=size,
max_tridiag_iter=10,
tolerance=0)
# Check cg
matrix_chol = batch_potrf(matrix)
actual = batch_potrs(rhs, matrix_chol)
self.assertTrue(approx_equal(solves, actual))
# Check tridiag
for i in range(5):
eigs = matrix[i].symeig()[0]
for j in range(8):
approx_eigs = t_mats[j, i].symeig()[0]
self.assertLess(
torch.mean(torch.abs((eigs - approx_eigs) / eigs)), 0.05)
def test_batch_cg(self):
batch = 5
size = 100
matrix = torch.DoubleTensor(batch, size, size).normal_()
matrix = matrix.matmul(matrix.transpose(-1, -2))
matrix.div_(matrix.norm())
matrix.add_(torch.DoubleTensor(matrix.size(-1)).fill_(1e-1).diag())
rhs = torch.DoubleTensor(batch, size, 50).normal_()
solves = linear_cg(matrix.matmul, rhs=rhs, max_iter=size)
# Check cg
matrix_chol = batch_potrf(matrix)
actual = batch_potrs(rhs, matrix_chol)
self.assertTrue(approx_equal(solves, actual))
def test_batch_cg(self):
batch = 5
size = 100
matrix = torch.randn(batch, size, size, dtype=torch.float64)
matrix = matrix.matmul(matrix.transpose(-1, -2))
matrix.div_(matrix.norm())
matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))
rhs = torch.randn(batch, size, 50, dtype=torch.float64)
solves = linear_cg(matrix.matmul, rhs=rhs, max_iter=size)
# Check cg
matrix_chol = batch_potrf(matrix)
actual = batch_potrs(rhs, matrix_chol)
self.assertTrue(approx_equal(solves, actual))