From 3fd065e98db60a847784c36a1e6b6dd9ce4d1f32 Mon Sep 17 00:00:00 2001 From: Kunal Ghosh Date: Tue, 2 Aug 2022 17:47:59 +0300 Subject: [PATCH 01/21] Initial commit of the linear conjugate gradients --- pymc_experimental/utils/linear_cg.py | 285 +++++++++++++++++++++++++++ 1 file changed, 285 insertions(+) create mode 100644 pymc_experimental/utils/linear_cg.py diff --git a/pymc_experimental/utils/linear_cg.py b/pymc_experimental/utils/linear_cg.py new file mode 100644 index 000000000..cc6befaac --- /dev/null +++ b/pymc_experimental/utils/linear_cg.py @@ -0,0 +1,285 @@ +from collections import namedtuple +import numpy as np + +pp = lambda x: np.array2string(x, precision=4, floatmode="fixed") +pp2 = lambda x: np.array2string(x, precision=2, floatmode="fixed") +Setting = namedtuple("setting", "on") + + +class settings: + eval_cg_tolerance = 0.01 + cg_tolerance = 1 + + def terminate_cg_by_size_on(self): + return False + + def _use_eval_tolerance_on(self): + return False + + terminate_cg_by_size = Setting(on=terminate_cg_by_size_on) + _use_eval_tolerance = Setting(on=_use_eval_tolerance_on) + + +print(f"SETTING : {settings.terminate_cg_by_size.on()}") + + +def masked_fill(vector, mask, fill_value): + masked_vector = np.ma.array(vector, mask=mask) + vector = masked_vector.filled(fill_value=fill_value) + return vector + + +def linear_cg_updates( + result, alpha, residual_inner_prod, eps, beta, residual, precond_residual, curr_conjugate_vec +): + + # Everything inside _jit_linear_cg_updates + result = result + alpha * curr_conjugate_vec + beta = np.copy(residual_inner_prod) + + residual_inner_prod = residual.T @ precond_residual + + # safe division + is_zero = beta < eps + beta = masked_fill(beta, mask=is_zero, fill_value=1) + + beta = residual_inner_prod / beta + beta = masked_fill(beta, mask=is_zero, fill_value=0) + curr_conjugate_vec = beta * curr_conjugate_vec + precond_residual + return ( + result, + alpha, + residual_inner_prod, + eps, + beta, + residual, + precond_residual, + curr_conjugate_vec, + ) + + +def linear_cg_np( + mat: np.matrix, + rhs, + n_tridiag=0, + tolerance=None, + eps=1e-10, + stop_updating_after=1e-10, + max_iter=1000, + max_tridiag_iter=20, + initial_guess=None, + preconditioner=None, +): + + if initial_guess is None: + initial_guess = np.zeros_like(rhs) + + if preconditioner is None: + preconditioner = lambda x: x + precond = False + else: + precond = True + + if tolerance is None: + if settings._use_eval_tolerance.on(): + tolerance = settings.eval_cg_tolerance + else: + tolerance = settings.cg_tolerance + + # If we are running m CG iterations, we obviously can't get more than m Lanczos coefficients + if max_tridiag_iter > max_iter: + raise RuntimeError( + "Getting a tridiagonalization larger than the number of CG iterations run is not possible!" + ) + + is_vector = len(rhs.shape) == 1 + if is_vector: + rhs = rhs[:, np.newaxis] + + num_rows = rhs.size + n_iter = min(max_iter, num_rows) if settings.terminate_cg_by_size.on() else max_iter + print(n_iter) + n_tridiag_iter = min(max_tridiag_iter, num_rows) + + # norm of rhs for convergence tests + rhs_norm = np.linalg.norm(rhs, 2) + # make almost-zero norms be 1 (so we don't get divide-by-zero errors) + rhs_is_zero = rhs_norm < eps + rhs_norm = masked_fill(rhs_norm, mask=rhs_is_zero, fill_value=1) + + # lets normalize rhs + rhs = rhs / rhs_norm + + # residuals + residual = rhs - mat @ initial_guess + batch_shape = residual.shape[:-2] + + result = np.copy(initial_guess) + + if np.allclose(residual, residual): + raise RuntimeError("NaNs encountered when trying to perform matrix-vector multiplication") + + # sometimes we are lucky and preconditioner solves the system right away + # check for convergence + residual_norm = np.linalg.norm(residual, 2) + has_converged = residual_norm < stop_updating_after + + if has_converged.all() and not n_tridiag: + n_iter = 0 # skip iterations + else: + precond_residual = preconditioner(residual) + + curr_conjugate_vec = precond_residual + residual_inner_prod = residual.T @ precond_residual + + # define storage matrices + mul_storage = np.zeros_like(residual) + alpha = np.zeros((*batch_shape, 1, rhs.shape[-1])) + beta = np.zeros_like(alpha) + is_zero = np.zeros((*batch_shape, 1, rhs.shape[-1])) + + # Define tridiagonal matrices if applicable + if n_tridiag: + t_mat = np.zeros((n_tridiag_iter, n_tridiag_iter, *batch_shape, n_tridiag)) + alpha_tridiag_is_zero = np.zeros(*batch_shape, n_tridiag) + alpha_reciprocal = np.zeros(*batch_shape, n_tridiag) + prev_alpha_reciprocal = np.zeros_like(alpha_reciprocal) + prev_beta = np.zeros_like(alpha_reciprocal) + + update_tridiag = True + last_tridiag_iter = 0 + + # it is possible that we don't reach tolerance even after all the iterations are over + tolerance_reached = False + + # start iteration + for k in range(n_iter): + mvms = mat @ curr_conjugate_vec + if precond: + alpha = curr_conjugate_vec @ mvms # scalar + + # safe division + is_zero = alpha < eps + alpha = masked_fill(alpha, mask=is_zero, fill_value=1) + alpha = residual_inner_prod / alpha + alpha = masked_fill(alpha, mask=is_zero, fill_value=0) + + # cancel out updates by setting directions which have converged to zero + alpha = masked_fill(alpha, mask=has_converged, fill_value=0) + residual = residual - alpha * mvms + + # update precond_residual + precond_residual = preconditioner(residual) + + # Everything inside _jit_linear_cg_updates + linear_cg_retvals = linear_cg_updates( + result, + alpha, + residual_inner_prod, + eps, + beta, + residual, + precond_residual, + curr_conjugate_vec, + ) + + ( + result, + alpha, + residual_inner_prod, + eps, + beta, + residual, + precond_residual, + curr_conjugate_vec, + ) = linear_cg_retvals + else: + # everything inside _jit_linear_cg_updates_no_precond + alpha = curr_conjugate_vec.T @ mvms + + # safe division + is_zero = alpha < eps + alpha = masked_fill(alpha, mask=is_zero, fill_value=1) + alpha = residual_inner_prod / alpha + alpha = masked_fill(alpha, is_zero, fill_value=0) + + alpha = masked_fill(alpha, has_converged, fill_value=0) # <- I'm here + residual = residual - alpha * mvms + precond_residual = np.copy(residual) + + linear_cg_retvals = linear_cg_updates( + result, + alpha, + residual_inner_prod, + eps, + beta, + residual, + precond_residual, + curr_conjugate_vec, + ) + ( + result, + alpha, + residual_inner_prod, + eps, + beta, + residual, + precond_residual, + curr_conjugate_vec, + ) = linear_cg_retvals + + residual_norm = np.linalg.norm(residual, 2) + residual_norm = masked_fill(residual_norm, mask=rhs_is_zero, fill_value=0) + has_converged = residual_norm < stop_updating_after + + if ( + k >= min(10, max_iter - 1) + and bool(residual_norm.mean() < tolerance) + and not (n_tridiag and k < min(n_tridiag_iter, max_iter - 1)) + ): + tolerance_reached = True + break + + # Update tridiagonal matrices, if applicable + if n_tridiag and k < n_tridiag_iter and update_tridiag: + alpha_tridiag = np.copy(alpha) + beta_tridiag = np.copy(beta) + + alpha_tridiag_is_zero = alpha_tridiag == 0 + alpha_tridiag = masked_fill(alpha_tridiag, mask=alpha_tridiag_is_zero, fill_value=1) + alpha_reciprocal = 1 / alpha_tridiag + alpha_tridiag = masked_fill(alpha_tridiag, mask=alpha_tridiag_is_zero, fill_value=0) + + if k == 0: + t_mat[k, k] = alpha_reciprocal + else: + t_mat[k, k] += np.squeeze(alpha_reciprocal + prev_beta * prev_alpha_reciprocal) + t_mat[k, k - 1] = np.sqrt(prev_beta) * prev_alpha_reciprocal + t_mat[k - 1, k] = np.copy(t_mat[k, k - 1]) + + if t_mat[k - 1, k].max() < 1e-6: + update_tridiag = False + + last_tridiag_iter = k + + prev_alpha_reciprocal = np.copy(alpha_reciprocal) + prev_beta = np.copy(beta_tridiag) + + # Un-normalize + result = result * rhs_norm + if not tolerance_reached and n_iter > 0: + raise RuntimeError( + "CG terminated in {} iterations with average residual norm {}" + " which is larger than the tolerance of {} specified by" + " gpytorch.settings.cg_tolerance." + " If performance is affected, consider raising the maximum number of CG iterations by running code in" + " a gpytorch.settings.max_cg_iterations(value) context.".format( + k + 1, residual_norm.mean(), tolerance + ) + ) + + if n_tridiag: + t_mat = t_mat[: last_tridiag_iter + 1, : last_tridiag_iter + 1] + return result, t_mat.transpose(-1, *range(2, 2 + len(batch_shape)), 0, 1) + else: + return result From e1cf869f6a06d00eb1b61527503ba12874bcf625 Mon Sep 17 00:00:00 2001 From: Kunal Ghosh Date: Tue, 2 Aug 2022 18:16:45 +0300 Subject: [PATCH 02/21] Updated __init__ so that we can import linear_cg --- pymc_experimental/utils/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc_experimental/utils/__init__.py b/pymc_experimental/utils/__init__.py index 6acbb2356..d155a5424 100644 --- a/pymc_experimental/utils/__init__.py +++ b/pymc_experimental/utils/__init__.py @@ -1 +1 @@ -from pymc_experimental.utils import prior, spline +from pymc_experimental.utils import prior, spline, linear_cg From 680afa6c5811f2834c76783dd60ac25ff6358013 Mon Sep 17 00:00:00 2001 From: Kunal Ghosh Date: Tue, 2 Aug 2022 18:20:25 +0300 Subject: [PATCH 03/21] Initial commit of pivoted cholesky --- pymc_experimental/utils/__init__.py | 2 +- pymc_experimental/utils/pivoted_cholesky.py | 63 +++++++++++++++++++++ 2 files changed, 64 insertions(+), 1 deletion(-) create mode 100644 pymc_experimental/utils/pivoted_cholesky.py diff --git a/pymc_experimental/utils/__init__.py b/pymc_experimental/utils/__init__.py index d155a5424..3c908356e 100644 --- a/pymc_experimental/utils/__init__.py +++ b/pymc_experimental/utils/__init__.py @@ -1 +1 @@ -from pymc_experimental.utils import prior, spline, linear_cg +from pymc_experimental.utils import prior, spline, linear_cg, pivoted_cholesky diff --git a/pymc_experimental/utils/pivoted_cholesky.py b/pymc_experimental/utils/pivoted_cholesky.py new file mode 100644 index 000000000..c6f286940 --- /dev/null +++ b/pymc_experimental/utils/pivoted_cholesky.py @@ -0,0 +1,63 @@ +import numpy as np +import torch + +from gpytorch.utils.permutation import apply_permutation + +pp = lambda x: np.array2string(x, precision=4, floatmode="fixed") + + +def pivoted_cholesky_np_gpt(mat: np.matrix, error_tol=1e-6, max_iter=np.Infinity): + """ + mat: numpy matrix of N x N + + This is to replicate what is done in GPyTorch verbatim. + """ + n = mat.shape[-1] + max_iter = min(int(max_iter), n) + + d = np.array(np.diag(mat)) + orig_error = np.max(d) + error = np.linalg.norm(d, 1) / orig_error + pi = np.arange(n) + + L = np.zeros((max_iter, n)) + + m = 0 + while m < max_iter and error > error_tol: + permuted_d = d[pi] + max_diag_idx = np.argmax(permuted_d[m:]) + max_diag_idx = max_diag_idx + m + max_diag_val = permuted_d[max_diag_idx] + i = max_diag_idx + + # swap pi_m and pi_i + pi[m], pi[i] = pi[i], pi[m] + pim = pi[m] + + L[m, pim] = np.sqrt(max_diag_val) + + if m + 1 < n: + row = apply_permutation( + torch.from_numpy(mat), torch.tensor(pim), right_permutation=None + ) # left permutation just swaps row + row = row.numpy().flatten() + pi_i = pi[m + 1 :] + L_m_new = row[pi_i] # length = 9 + + if m > 0: + L_prev = L[:m, pi_i] + update = L[:m, pim] + prod = update @ L_prev + L_m_new = L_m_new - prod # np.sum(prod, axis=-1) + + L_m = L[m, :] + L_m_new = L_m_new / L_m[pim] + L_m[pi_i] = L_m_new + + matrix_diag_current = d[pi_i] + d[pi_i] = matrix_diag_current - L_m_new**2 + + L[m, :] = L_m + error = np.linalg.norm(d[pi_i], 1) / orig_error + m = m + 1 + return L, pi From c42bb358a8f0c11e0a5084cc90eea4192a1f3bb9 Mon Sep 17 00:00:00 2001 From: Kunal Ghosh Date: Tue, 2 Aug 2022 18:23:18 +0300 Subject: [PATCH 04/21] Fixed the name of pivoted cholesky function --- pymc_experimental/utils/pivoted_cholesky.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc_experimental/utils/pivoted_cholesky.py b/pymc_experimental/utils/pivoted_cholesky.py index c6f286940..ecb1106e1 100644 --- a/pymc_experimental/utils/pivoted_cholesky.py +++ b/pymc_experimental/utils/pivoted_cholesky.py @@ -6,7 +6,7 @@ pp = lambda x: np.array2string(x, precision=4, floatmode="fixed") -def pivoted_cholesky_np_gpt(mat: np.matrix, error_tol=1e-6, max_iter=np.Infinity): +def pivoted_cholesky(mat: np.matrix, error_tol=1e-6, max_iter=np.Infinity): """ mat: numpy matrix of N x N From 21ed2d5ec67b6a944ba8382978275a526566b5f0 Mon Sep 17 00:00:00 2001 From: Kunal Ghosh Date: Tue, 2 Aug 2022 18:42:04 +0300 Subject: [PATCH 05/21] Since we are invoking the function as a class attribute, removing the self --- pymc_experimental/utils/linear_cg.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pymc_experimental/utils/linear_cg.py b/pymc_experimental/utils/linear_cg.py index cc6befaac..082d51c34 100644 --- a/pymc_experimental/utils/linear_cg.py +++ b/pymc_experimental/utils/linear_cg.py @@ -10,10 +10,10 @@ class settings: eval_cg_tolerance = 0.01 cg_tolerance = 1 - def terminate_cg_by_size_on(self): + def terminate_cg_by_size_on(): return False - def _use_eval_tolerance_on(self): + def _use_eval_tolerance_on(): return False terminate_cg_by_size = Setting(on=terminate_cg_by_size_on) From 021a3e7e02baf9d40838e8427409f2a873515b2e Mon Sep 17 00:00:00 2001 From: Kunal Ghosh Date: Tue, 2 Aug 2022 18:46:35 +0300 Subject: [PATCH 06/21] Removing linear conjugate gradients from this branch --- pymc_experimental/utils/linear_cg.py | 285 --------------------------- 1 file changed, 285 deletions(-) delete mode 100644 pymc_experimental/utils/linear_cg.py diff --git a/pymc_experimental/utils/linear_cg.py b/pymc_experimental/utils/linear_cg.py deleted file mode 100644 index 082d51c34..000000000 --- a/pymc_experimental/utils/linear_cg.py +++ /dev/null @@ -1,285 +0,0 @@ -from collections import namedtuple -import numpy as np - -pp = lambda x: np.array2string(x, precision=4, floatmode="fixed") -pp2 = lambda x: np.array2string(x, precision=2, floatmode="fixed") -Setting = namedtuple("setting", "on") - - -class settings: - eval_cg_tolerance = 0.01 - cg_tolerance = 1 - - def terminate_cg_by_size_on(): - return False - - def _use_eval_tolerance_on(): - return False - - terminate_cg_by_size = Setting(on=terminate_cg_by_size_on) - _use_eval_tolerance = Setting(on=_use_eval_tolerance_on) - - -print(f"SETTING : {settings.terminate_cg_by_size.on()}") - - -def masked_fill(vector, mask, fill_value): - masked_vector = np.ma.array(vector, mask=mask) - vector = masked_vector.filled(fill_value=fill_value) - return vector - - -def linear_cg_updates( - result, alpha, residual_inner_prod, eps, beta, residual, precond_residual, curr_conjugate_vec -): - - # Everything inside _jit_linear_cg_updates - result = result + alpha * curr_conjugate_vec - beta = np.copy(residual_inner_prod) - - residual_inner_prod = residual.T @ precond_residual - - # safe division - is_zero = beta < eps - beta = masked_fill(beta, mask=is_zero, fill_value=1) - - beta = residual_inner_prod / beta - beta = masked_fill(beta, mask=is_zero, fill_value=0) - curr_conjugate_vec = beta * curr_conjugate_vec + precond_residual - return ( - result, - alpha, - residual_inner_prod, - eps, - beta, - residual, - precond_residual, - curr_conjugate_vec, - ) - - -def linear_cg_np( - mat: np.matrix, - rhs, - n_tridiag=0, - tolerance=None, - eps=1e-10, - stop_updating_after=1e-10, - max_iter=1000, - max_tridiag_iter=20, - initial_guess=None, - preconditioner=None, -): - - if initial_guess is None: - initial_guess = np.zeros_like(rhs) - - if preconditioner is None: - preconditioner = lambda x: x - precond = False - else: - precond = True - - if tolerance is None: - if settings._use_eval_tolerance.on(): - tolerance = settings.eval_cg_tolerance - else: - tolerance = settings.cg_tolerance - - # If we are running m CG iterations, we obviously can't get more than m Lanczos coefficients - if max_tridiag_iter > max_iter: - raise RuntimeError( - "Getting a tridiagonalization larger than the number of CG iterations run is not possible!" - ) - - is_vector = len(rhs.shape) == 1 - if is_vector: - rhs = rhs[:, np.newaxis] - - num_rows = rhs.size - n_iter = min(max_iter, num_rows) if settings.terminate_cg_by_size.on() else max_iter - print(n_iter) - n_tridiag_iter = min(max_tridiag_iter, num_rows) - - # norm of rhs for convergence tests - rhs_norm = np.linalg.norm(rhs, 2) - # make almost-zero norms be 1 (so we don't get divide-by-zero errors) - rhs_is_zero = rhs_norm < eps - rhs_norm = masked_fill(rhs_norm, mask=rhs_is_zero, fill_value=1) - - # lets normalize rhs - rhs = rhs / rhs_norm - - # residuals - residual = rhs - mat @ initial_guess - batch_shape = residual.shape[:-2] - - result = np.copy(initial_guess) - - if np.allclose(residual, residual): - raise RuntimeError("NaNs encountered when trying to perform matrix-vector multiplication") - - # sometimes we are lucky and preconditioner solves the system right away - # check for convergence - residual_norm = np.linalg.norm(residual, 2) - has_converged = residual_norm < stop_updating_after - - if has_converged.all() and not n_tridiag: - n_iter = 0 # skip iterations - else: - precond_residual = preconditioner(residual) - - curr_conjugate_vec = precond_residual - residual_inner_prod = residual.T @ precond_residual - - # define storage matrices - mul_storage = np.zeros_like(residual) - alpha = np.zeros((*batch_shape, 1, rhs.shape[-1])) - beta = np.zeros_like(alpha) - is_zero = np.zeros((*batch_shape, 1, rhs.shape[-1])) - - # Define tridiagonal matrices if applicable - if n_tridiag: - t_mat = np.zeros((n_tridiag_iter, n_tridiag_iter, *batch_shape, n_tridiag)) - alpha_tridiag_is_zero = np.zeros(*batch_shape, n_tridiag) - alpha_reciprocal = np.zeros(*batch_shape, n_tridiag) - prev_alpha_reciprocal = np.zeros_like(alpha_reciprocal) - prev_beta = np.zeros_like(alpha_reciprocal) - - update_tridiag = True - last_tridiag_iter = 0 - - # it is possible that we don't reach tolerance even after all the iterations are over - tolerance_reached = False - - # start iteration - for k in range(n_iter): - mvms = mat @ curr_conjugate_vec - if precond: - alpha = curr_conjugate_vec @ mvms # scalar - - # safe division - is_zero = alpha < eps - alpha = masked_fill(alpha, mask=is_zero, fill_value=1) - alpha = residual_inner_prod / alpha - alpha = masked_fill(alpha, mask=is_zero, fill_value=0) - - # cancel out updates by setting directions which have converged to zero - alpha = masked_fill(alpha, mask=has_converged, fill_value=0) - residual = residual - alpha * mvms - - # update precond_residual - precond_residual = preconditioner(residual) - - # Everything inside _jit_linear_cg_updates - linear_cg_retvals = linear_cg_updates( - result, - alpha, - residual_inner_prod, - eps, - beta, - residual, - precond_residual, - curr_conjugate_vec, - ) - - ( - result, - alpha, - residual_inner_prod, - eps, - beta, - residual, - precond_residual, - curr_conjugate_vec, - ) = linear_cg_retvals - else: - # everything inside _jit_linear_cg_updates_no_precond - alpha = curr_conjugate_vec.T @ mvms - - # safe division - is_zero = alpha < eps - alpha = masked_fill(alpha, mask=is_zero, fill_value=1) - alpha = residual_inner_prod / alpha - alpha = masked_fill(alpha, is_zero, fill_value=0) - - alpha = masked_fill(alpha, has_converged, fill_value=0) # <- I'm here - residual = residual - alpha * mvms - precond_residual = np.copy(residual) - - linear_cg_retvals = linear_cg_updates( - result, - alpha, - residual_inner_prod, - eps, - beta, - residual, - precond_residual, - curr_conjugate_vec, - ) - ( - result, - alpha, - residual_inner_prod, - eps, - beta, - residual, - precond_residual, - curr_conjugate_vec, - ) = linear_cg_retvals - - residual_norm = np.linalg.norm(residual, 2) - residual_norm = masked_fill(residual_norm, mask=rhs_is_zero, fill_value=0) - has_converged = residual_norm < stop_updating_after - - if ( - k >= min(10, max_iter - 1) - and bool(residual_norm.mean() < tolerance) - and not (n_tridiag and k < min(n_tridiag_iter, max_iter - 1)) - ): - tolerance_reached = True - break - - # Update tridiagonal matrices, if applicable - if n_tridiag and k < n_tridiag_iter and update_tridiag: - alpha_tridiag = np.copy(alpha) - beta_tridiag = np.copy(beta) - - alpha_tridiag_is_zero = alpha_tridiag == 0 - alpha_tridiag = masked_fill(alpha_tridiag, mask=alpha_tridiag_is_zero, fill_value=1) - alpha_reciprocal = 1 / alpha_tridiag - alpha_tridiag = masked_fill(alpha_tridiag, mask=alpha_tridiag_is_zero, fill_value=0) - - if k == 0: - t_mat[k, k] = alpha_reciprocal - else: - t_mat[k, k] += np.squeeze(alpha_reciprocal + prev_beta * prev_alpha_reciprocal) - t_mat[k, k - 1] = np.sqrt(prev_beta) * prev_alpha_reciprocal - t_mat[k - 1, k] = np.copy(t_mat[k, k - 1]) - - if t_mat[k - 1, k].max() < 1e-6: - update_tridiag = False - - last_tridiag_iter = k - - prev_alpha_reciprocal = np.copy(alpha_reciprocal) - prev_beta = np.copy(beta_tridiag) - - # Un-normalize - result = result * rhs_norm - if not tolerance_reached and n_iter > 0: - raise RuntimeError( - "CG terminated in {} iterations with average residual norm {}" - " which is larger than the tolerance of {} specified by" - " gpytorch.settings.cg_tolerance." - " If performance is affected, consider raising the maximum number of CG iterations by running code in" - " a gpytorch.settings.max_cg_iterations(value) context.".format( - k + 1, residual_norm.mean(), tolerance - ) - ) - - if n_tridiag: - t_mat = t_mat[: last_tridiag_iter + 1, : last_tridiag_iter + 1] - return result, t_mat.transpose(-1, *range(2, 2 + len(batch_shape)), 0, 1) - else: - return result From bbe525bef9fb82e45755f69538505f974dc149ab Mon Sep 17 00:00:00 2001 From: Kunal Ghosh Date: Tue, 2 Aug 2022 18:47:36 +0300 Subject: [PATCH 07/21] Also removing linear cg from __init__ --- pymc_experimental/utils/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc_experimental/utils/__init__.py b/pymc_experimental/utils/__init__.py index 3c908356e..d3a2d616a 100644 --- a/pymc_experimental/utils/__init__.py +++ b/pymc_experimental/utils/__init__.py @@ -1 +1 @@ -from pymc_experimental.utils import prior, spline, linear_cg, pivoted_cholesky +from pymc_experimental.utils import prior, spline, pivoted_cholesky From 9ed06f3a5c604e15941ab5beb90842eb6a711f77 Mon Sep 17 00:00:00 2001 From: Kunal Ghosh Date: Tue, 2 Aug 2022 18:53:44 +0300 Subject: [PATCH 08/21] Fixed the import --- pymc_experimental/utils/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pymc_experimental/utils/__init__.py b/pymc_experimental/utils/__init__.py index d3a2d616a..e44211981 100644 --- a/pymc_experimental/utils/__init__.py +++ b/pymc_experimental/utils/__init__.py @@ -1 +1,2 @@ -from pymc_experimental.utils import prior, spline, pivoted_cholesky +from pymc_experimental.utils import prior, spline +from pymc_experimental.utils.pivoted_cholesky import pivoted_cholesky From 2d5e052402fec143772688e3991a1070420669bc Mon Sep 17 00:00:00 2001 From: Kunal Ghosh Date: Wed, 3 Aug 2022 14:37:11 +0300 Subject: [PATCH 09/21] Adding dependencies for pivoted_cholesky, they are also needed for tests --- requirements.txt | 2 ++ 1 file changed, 2 insertions(+) diff --git a/requirements.txt b/requirements.txt index c40987ae7..0dba2bcbc 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,4 @@ pymc>=4.0.1 xhistogram +gpytorch +pytorch From cd957c81717f9f44c0eabcd685da0086b471d674 Mon Sep 17 00:00:00 2001 From: Kunal Ghosh Date: Thu, 11 Aug 2022 13:53:36 +0300 Subject: [PATCH 10/21] Added correct package for pytorch --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 0dba2bcbc..0f851eef9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ pymc>=4.0.1 xhistogram gpytorch -pytorch +torch From f0b48555f256330d5e60d7325fd42a479d1a174b Mon Sep 17 00:00:00 2001 From: Kunal Ghosh Date: Tue, 27 Sep 2022 20:23:49 +0300 Subject: [PATCH 11/21] Added try...catch block to ensure users who don't have these packages installed, and don't want to use pivoted Cholesky, can use pymc --- pymc_experimental/utils/pivoted_cholesky.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/pymc_experimental/utils/pivoted_cholesky.py b/pymc_experimental/utils/pivoted_cholesky.py index ecb1106e1..b626d60c5 100644 --- a/pymc_experimental/utils/pivoted_cholesky.py +++ b/pymc_experimental/utils/pivoted_cholesky.py @@ -1,7 +1,11 @@ +try: + from gpytorch.utils.permutation import apply_permutation + import torch +except ImportError as e: + raise ImportError( + "Please install Pytorch and GPyTorch to use this pivoted Cholesky implementation." + ) import numpy as np -import torch - -from gpytorch.utils.permutation import apply_permutation pp = lambda x: np.array2string(x, precision=4, floatmode="fixed") From f3a13f5d325420249401f028b6b7edd9ee2a6064 Mon Sep 17 00:00:00 2001 From: Kunal Ghosh Date: Tue, 27 Sep 2022 20:35:26 +0300 Subject: [PATCH 12/21] Added import checks in the test file as well --- .../tests/test_pivoted_cholesky.py | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) create mode 100644 pymc_experimental/tests/test_pivoted_cholesky.py diff --git a/pymc_experimental/tests/test_pivoted_cholesky.py b/pymc_experimental/tests/test_pivoted_cholesky.py new file mode 100644 index 000000000..3fdd853d3 --- /dev/null +++ b/pymc_experimental/tests/test_pivoted_cholesky.py @@ -0,0 +1,23 @@ +try: + import torch + import gpytorch +except ImportError as e: + raise ImportError( + "Please install Pytorch and GPyTorch to use this pivoted Cholesky implementation." + ) +import numpy as np +import pytest +import pymc_experimental as pmx + + +def test_match_gpytorch_linearcg_output(): + N = 10 + rank = 5 + np.random.seed(1234) # nans with seed 1234 + K = np.random.randn(N, N) + K = K @ K.T + N * np.eye(N) + K_torch = torch.from_numpy(K) + + L_gpt = gpytorch.pivoted_cholesky(K_torch, rank=rank, error_tol=1e-3) + L_np, _ = pmx.utils.pivoted_cholesky(K, max_iter=rank, error_tol=1e-3) + assert np.allclose(L_gpt, L_np.T) From 03b7bf7396f48bb41e51e5cb83509dc6163fba25 Mon Sep 17 00:00:00 2001 From: Kunal Ghosh Date: Tue, 27 Sep 2022 21:30:41 +0300 Subject: [PATCH 13/21] removed unused commits --- pymc_experimental/tests/test_pivoted_cholesky.py | 4 ++-- pymc_experimental/utils/pivoted_cholesky.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pymc_experimental/tests/test_pivoted_cholesky.py b/pymc_experimental/tests/test_pivoted_cholesky.py index 3fdd853d3..e4458cfae 100644 --- a/pymc_experimental/tests/test_pivoted_cholesky.py +++ b/pymc_experimental/tests/test_pivoted_cholesky.py @@ -1,12 +1,12 @@ try: - import torch import gpytorch + import torch except ImportError as e: raise ImportError( "Please install Pytorch and GPyTorch to use this pivoted Cholesky implementation." ) import numpy as np -import pytest + import pymc_experimental as pmx diff --git a/pymc_experimental/utils/pivoted_cholesky.py b/pymc_experimental/utils/pivoted_cholesky.py index b626d60c5..a125bf69f 100644 --- a/pymc_experimental/utils/pivoted_cholesky.py +++ b/pymc_experimental/utils/pivoted_cholesky.py @@ -1,6 +1,6 @@ try: - from gpytorch.utils.permutation import apply_permutation import torch + from gpytorch.utils.permutation import apply_permutation except ImportError as e: raise ImportError( "Please install Pytorch and GPyTorch to use this pivoted Cholesky implementation." From 90094276849b388e9353e5eec1e2dd8c5317d23c Mon Sep 17 00:00:00 2001 From: Kunal Ghosh Date: Tue, 27 Sep 2022 21:56:18 +0300 Subject: [PATCH 14/21] removing pytorch and gpytorch from requirements. --- requirements.txt | 2 -- 1 file changed, 2 deletions(-) diff --git a/requirements.txt b/requirements.txt index 0f851eef9..c40987ae7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,2 @@ pymc>=4.0.1 xhistogram -gpytorch -torch From 852d95cd9b22d117816eee72fa910e1cc82bdf2d Mon Sep 17 00:00:00 2001 From: Kunal Ghosh Date: Tue, 27 Sep 2022 22:13:05 +0300 Subject: [PATCH 15/21] pre-commit wouldn't let me commit print statements --- pymc_experimental/tests/test_pivoted_cholesky.py | 7 ++++--- pymc_experimental/utils/pivoted_cholesky.py | 7 ++++--- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/pymc_experimental/tests/test_pivoted_cholesky.py b/pymc_experimental/tests/test_pivoted_cholesky.py index e4458cfae..72fd045c3 100644 --- a/pymc_experimental/tests/test_pivoted_cholesky.py +++ b/pymc_experimental/tests/test_pivoted_cholesky.py @@ -2,9 +2,10 @@ import gpytorch import torch except ImportError as e: - raise ImportError( - "Please install Pytorch and GPyTorch to use this pivoted Cholesky implementation." - ) + # print( + # f"Please install Pytorch and GPyTorch to use this pivoted Cholesky implementation. Error {e}" + # ) + pass import numpy as np import pymc_experimental as pmx diff --git a/pymc_experimental/utils/pivoted_cholesky.py b/pymc_experimental/utils/pivoted_cholesky.py index a125bf69f..3e61f780f 100644 --- a/pymc_experimental/utils/pivoted_cholesky.py +++ b/pymc_experimental/utils/pivoted_cholesky.py @@ -2,9 +2,10 @@ import torch from gpytorch.utils.permutation import apply_permutation except ImportError as e: - raise ImportError( - "Please install Pytorch and GPyTorch to use this pivoted Cholesky implementation." - ) + # print( + # f"Please install Pytorch and GPyTorch to use this pivoted Cholesky implementation. Error {e}" + # ) + pass import numpy as np pp = lambda x: np.array2string(x, precision=4, floatmode="fixed") From 7d1fb4e0cace825aaee5c225bbc5b057dafcc311 Mon Sep 17 00:00:00 2001 From: Kunal Ghosh Date: Tue, 27 Sep 2022 22:24:00 +0300 Subject: [PATCH 16/21] removing the test for now --- .../tests/test_pivoted_cholesky.py | 48 +++++++++---------- 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/pymc_experimental/tests/test_pivoted_cholesky.py b/pymc_experimental/tests/test_pivoted_cholesky.py index 72fd045c3..88ad75ce5 100644 --- a/pymc_experimental/tests/test_pivoted_cholesky.py +++ b/pymc_experimental/tests/test_pivoted_cholesky.py @@ -1,24 +1,24 @@ -try: - import gpytorch - import torch -except ImportError as e: - # print( - # f"Please install Pytorch and GPyTorch to use this pivoted Cholesky implementation. Error {e}" - # ) - pass -import numpy as np - -import pymc_experimental as pmx - - -def test_match_gpytorch_linearcg_output(): - N = 10 - rank = 5 - np.random.seed(1234) # nans with seed 1234 - K = np.random.randn(N, N) - K = K @ K.T + N * np.eye(N) - K_torch = torch.from_numpy(K) - - L_gpt = gpytorch.pivoted_cholesky(K_torch, rank=rank, error_tol=1e-3) - L_np, _ = pmx.utils.pivoted_cholesky(K, max_iter=rank, error_tol=1e-3) - assert np.allclose(L_gpt, L_np.T) +# try: +# import gpytorch +# import torch +# except ImportError as e: +# # print( +# # f"Please install Pytorch and GPyTorch to use this pivoted Cholesky implementation. Error {e}" +# # ) +# pass +# import numpy as np +# +# import pymc_experimental as pmx +# +# +# def test_match_gpytorch_linearcg_output(): +# N = 10 +# rank = 5 +# np.random.seed(1234) # nans with seed 1234 +# K = np.random.randn(N, N) +# K = K @ K.T + N * np.eye(N) +# K_torch = torch.from_numpy(K) +# +# L_gpt = gpytorch.pivoted_cholesky(K_torch, rank=rank, error_tol=1e-3) +# L_np, _ = pmx.utils.pivoted_cholesky(K, max_iter=rank, error_tol=1e-3) +# assert np.allclose(L_gpt, L_np.T) From 8aba7b0706d8fc31b6cc0ca205a0f877d4775b37 Mon Sep 17 00:00:00 2001 From: Kunal Ghosh Date: Fri, 14 Oct 2022 11:50:47 +0300 Subject: [PATCH 17/21] Raising an ImportError instead of printing --- pymc_experimental/utils/pivoted_cholesky.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/pymc_experimental/utils/pivoted_cholesky.py b/pymc_experimental/utils/pivoted_cholesky.py index 3e61f780f..c83ab3f53 100644 --- a/pymc_experimental/utils/pivoted_cholesky.py +++ b/pymc_experimental/utils/pivoted_cholesky.py @@ -2,10 +2,8 @@ import torch from gpytorch.utils.permutation import apply_permutation except ImportError as e: - # print( - # f"Please install Pytorch and GPyTorch to use this pivoted Cholesky implementation. Error {e}" - # ) - pass + raise ImportError("PyTorch and GPyTorch not found.") from e + import numpy as np pp = lambda x: np.array2string(x, precision=4, floatmode="fixed") From 9b8a13291a54c725142017fbde34d233baea2af6 Mon Sep 17 00:00:00 2001 From: Kunal Ghosh Date: Fri, 14 Oct 2022 11:51:28 +0300 Subject: [PATCH 18/21] Addressing stylistic comment --- pymc_experimental/utils/pivoted_cholesky.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc_experimental/utils/pivoted_cholesky.py b/pymc_experimental/utils/pivoted_cholesky.py index c83ab3f53..a839e14ce 100644 --- a/pymc_experimental/utils/pivoted_cholesky.py +++ b/pymc_experimental/utils/pivoted_cholesky.py @@ -9,7 +9,7 @@ pp = lambda x: np.array2string(x, precision=4, floatmode="fixed") -def pivoted_cholesky(mat: np.matrix, error_tol=1e-6, max_iter=np.Infinity): +def pivoted_cholesky(mat: np.matrix, error_tol=1e-6, max_iter=np.inf): """ mat: numpy matrix of N x N From 3676a51d14999f6578b24edb4b7dbf375d0a619d Mon Sep 17 00:00:00 2001 From: Kunal Ghosh Date: Fri, 14 Oct 2022 12:02:42 +0300 Subject: [PATCH 19/21] formatting fixes --- pymc_experimental/utils/pivoted_cholesky.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc_experimental/utils/pivoted_cholesky.py b/pymc_experimental/utils/pivoted_cholesky.py index a839e14ce..69ea9cd7b 100644 --- a/pymc_experimental/utils/pivoted_cholesky.py +++ b/pymc_experimental/utils/pivoted_cholesky.py @@ -2,7 +2,7 @@ import torch from gpytorch.utils.permutation import apply_permutation except ImportError as e: - raise ImportError("PyTorch and GPyTorch not found.") from e + raise ImportError("PyTorch and GPyTorch not found.") from e import numpy as np From 4a9a8c06df225ccea0bd8ce82e6edea68ad2ce52 Mon Sep 17 00:00:00 2001 From: Kunal Ghosh Date: Fri, 14 Oct 2022 12:03:17 +0300 Subject: [PATCH 20/21] formatting fixes --- pymc_experimental/utils/__init__.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pymc_experimental/utils/__init__.py b/pymc_experimental/utils/__init__.py index 11b9646ee..a98f70d62 100644 --- a/pymc_experimental/utils/__init__.py +++ b/pymc_experimental/utils/__init__.py @@ -14,6 +14,5 @@ from pymc_experimental.utils import prior, spline - -from pymc_experimental.utils.pivoted_cholesky import pivoted_cholesky from pymc_experimental.utils.linear_cg import linear_cg +from pymc_experimental.utils.pivoted_cholesky import pivoted_cholesky From 96229a87868ca41827789ee6310ab6461b2a1d3d Mon Sep 17 00:00:00 2001 From: Kunal Ghosh Date: Mon, 17 Oct 2022 09:48:51 +0300 Subject: [PATCH 21/21] pre-commit modifications --- pymc_experimental/utils/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pymc_experimental/utils/__init__.py b/pymc_experimental/utils/__init__.py index a98f70d62..db751aa2a 100644 --- a/pymc_experimental/utils/__init__.py +++ b/pymc_experimental/utils/__init__.py @@ -15,4 +15,5 @@ from pymc_experimental.utils import prior, spline from pymc_experimental.utils.linear_cg import linear_cg -from pymc_experimental.utils.pivoted_cholesky import pivoted_cholesky + +# from pymc_experimental.utils.pivoted_cholesky import pivoted_cholesky