A few tweaks to linear_cg
The modified check for convergence is a little cleaner and also slightly faster.
@@ -79,7 +79,7 @@ def linear_cg(
# Sometime we're lucky and the preconditioner solves the system right away
residual_norm = residual.norm(2, dim=-2)
- if not torch.sum(residual_norm > tolerance) and not n_tridiag:
+ if torch.all(residual_norm < tolerance) and not n_tridiag:
n_iter = 0 # Skip the iteration!
# Otherwise, let's define precond_residual and curr_conjugate_vec
@@ -91,10 +91,10 @@ def linear_cg(
# Define storage matrices
mul_storage = residual.new(residual.size())
- alpha = residual.new(
- rhs.size(0), 1, rhs.size(-1)
- ) if rhs.ndimension() == 3 else residual.new(
- 1, rhs.size(-1)
+ alpha = (
+ residual.new(rhs.size(0), 1, rhs.size(-1))
+ if rhs.ndimension() == 3
+ else residual.new(1, rhs.size(-1))
)
beta = alpha.new(alpha.size())
@@ -131,7 +131,7 @@ def linear_cg(
# If residual are sufficiently small, then exit loop
# Alternatively, exit if this is our last iteration
torch.norm(residual, 2, dim=-2, out=residual_norm)
- if not (torch.sum(residual_norm > tolerance)) and not n_tridiag:
+ if torch.all(residual_norm < tolerance) and not n_tridiag:
break
# Update precond_residual
