From be0e68dbf835dcf7a8ea684ae88f63702fd84c92 Mon Sep 17 00:00:00 2001 From: Fritz Obermeyer Date: Thu, 27 Oct 2022 13:25:07 -0700 Subject: [PATCH 01/11] Attempt to add jitter to Gaussian ops --- pyro/distributions/hmm.py | 43 ++++++++++++++++------- pyro/ops/gaussian.py | 49 +++++++++++++++++++++------ pyro/ops/tensor_utils.py | 4 ++- tests/distributions/test_hmm.py | 3 +- tests/ops/test_gaussian.py | 60 ++++++++++++++++++++++++++++++++- 5 files changed, 133 insertions(+), 26 deletions(-) diff --git a/pyro/distributions/hmm.py b/pyro/distributions/hmm.py index 2a38b19f11..50162e4d29 100644 --- a/pyro/distributions/hmm.py +++ b/pyro/distributions/hmm.py @@ -489,6 +489,8 @@ class GaussianHMM(HiddenMarkovModel): :param int duration: Optional size of the time axis ``event_shape[0]``. This is required when sampling from homogeneous HMMs whose parameters are not expanded along the time axis. + :param float jitter: Optional constant added to matrix diagonals before + performing Cholesky decompositions, to improve stability. """ has_rsample = True @@ -504,6 +506,8 @@ def __init__( observation_dist, validate_args=None, duration=None, + *, + jitter: float = 0.0, ): assert isinstance(initial_dist, torch.distributions.MultivariateNormal) or ( isinstance(initial_dist, torch.distributions.Independent) @@ -542,11 +546,13 @@ def __init__( self._init = mvn_to_gaussian(initial_dist).expand(self.batch_shape) self._trans = matrix_and_mvn_to_gaussian(transition_matrix, transition_dist) self._obs = matrix_and_mvn_to_gaussian(observation_matrix, observation_dist) + self.jitter = jitter def expand(self, batch_shape, _instance=None): new = self._get_checked_instance(GaussianHMM, _instance) new.hidden_dim = self.hidden_dim new.obs_dim = self.obs_dim + new.jitter = self.jitter new._obs = self._obs new._trans = self._trans @@ -572,10 +578,14 @@ def log_prob(self, value): ) # Eliminate time dimension. - result = sequential_gaussian_tensordot(result.expand(result.batch_shape)) + result = sequential_gaussian_tensordot( + result.expand(result.batch_shape), jitter=self.jitter + ) # Combine initial factor. - result = gaussian_tensordot(self._init, result, dims=self.hidden_dim) + result = gaussian_tensordot( + self._init, result, dims=self.hidden_dim, jitter=self.jitter + ) # Marginalize out final state. result = result.event_logsumexp() @@ -584,11 +594,13 @@ def log_prob(self, value): def rsample(self, sample_shape=torch.Size()): assert self.duration is not None sample_shape = torch.Size(sample_shape) - trans = self._trans + self._obs.marginalize(right=self.obs_dim).event_pad( - left=self.hidden_dim - ) + trans = self._trans + self._obs.marginalize( + right=self.obs_dim, jitter=self.jitter + ).event_pad(left=self.hidden_dim) trans = trans.expand(trans.batch_shape[:-1] + (self.duration,)) - z = sequential_gaussian_filter_sample(self._init, trans, sample_shape) + z = sequential_gaussian_filter_sample( + self._init, trans, sample_shape, jitter=self.jitter + ) z = z[..., 1:, :] # drop the initial hidden state x = self._obs.left_condition(z).rsample() return x @@ -599,7 +611,9 @@ def rsample_posterior(self, value, sample_shape=torch.Size()): """ trans = self._trans + self._obs.condition(value).event_pad(left=self.hidden_dim) trans = trans.expand(trans.batch_shape) - z = sequential_gaussian_filter_sample(self._init, trans, sample_shape) + z = sequential_gaussian_filter_sample( + self._init, trans, sample_shape, jitter=self.jitter + ) z = z[..., 1:, :] # drop the initial hidden state return z @@ -621,10 +635,14 @@ def filter(self, value): logp = self._trans + self._obs.condition(value).event_pad(left=self.hidden_dim) # Eliminate time dimension. - logp = sequential_gaussian_tensordot(logp.expand(logp.batch_shape)) + logp = sequential_gaussian_tensordot( + logp.expand(logp.batch_shape), jitter=self.jitter + ) # Combine initial factor. - logp = gaussian_tensordot(self._init, logp, dims=self.hidden_dim) + logp = gaussian_tensordot( + self._init, logp, dims=self.hidden_dim, jitter=self.jitter + ) # Convert to a distribution precision = logp.precision @@ -664,6 +682,7 @@ def conjugate_update(self, other): new = self._get_checked_instance(GaussianHMM) new.hidden_dim = self.hidden_dim new.obs_dim = self.obs_dim + new.jitter = self.jitter new._init = self._init new._trans = self._trans new._obs = self._obs + mvn_to_gaussian(other.to_event(-1)).event_pad( @@ -672,9 +691,9 @@ def conjugate_update(self, other): # Normalize. # TODO cache this computation for the forward pass of .rsample(). - logp = new._trans + new._obs.marginalize(right=new.obs_dim).event_pad( - left=new.hidden_dim - ) + logp = new._trans + new._obs.marginalize( + right=new.obs_dim, jitter=self.jitter + ).event_pad(left=new.hidden_dim) logp = sequential_gaussian_tensordot(logp.expand(logp.batch_shape)) logp = gaussian_tensordot(new._init, logp, dims=new.hidden_dim) log_normalizer = logp.event_logsumexp() diff --git a/pyro/ops/gaussian.py b/pyro/ops/gaussian.py index 4541513495..fba0ec7a1f 100644 --- a/pyro/ops/gaussian.py +++ b/pyro/ops/gaussian.py @@ -9,7 +9,13 @@ from torch.nn.functional import pad from pyro.distributions.util import broadcast_shape -from pyro.ops.tensor_utils import cholesky, matmul, matvecmul, triangular_solve +from pyro.ops.tensor_utils import ( + cholesky, + cholesky_solve, + matmul, + matvecmul, + triangular_solve, +) class Gaussian: @@ -155,7 +161,7 @@ def rsample( Reparameterized sampler. """ P_chol = cholesky(self.precision) - loc = self.info_vec.unsqueeze(-1).cholesky_solve(P_chol).squeeze(-1) + loc = cholesky_solve(self.info_vec.unsqueeze(-1), P_chol).squeeze(-1) shape = sample_shape + self.batch_shape + (self.dim(), 1) if noise is None: noise = torch.randn(shape, dtype=loc.dtype, device=loc.device) @@ -230,7 +236,9 @@ def left_condition(self, value: torch.Tensor) -> "Gaussian": ) return self.event_permute(perm).condition(value) - def marginalize(self, left=0, right=0) -> "Gaussian": + def marginalize( + self, left: int = 0, right: int = 0, *, jitter: float = 0.0 + ) -> "Gaussian": """ Marginalizing out variables on either side of the event dimension:: @@ -241,6 +249,11 @@ def marginalize(self, left=0, right=0) -> "Gaussian": g.condition(x).event_logsumexp() = g.marginalize(left=g.dim() - x.size(-1)).log_density(x) + + :param int left: Number of left event dims to marginalize. + :param int right: Number of right event dims to marginalize. + :param float jitter: Optional constant added to matrix diagonals before + performing Cholesky decompositions, to improve stability. """ if left == 0 and right == 0: return self @@ -254,7 +267,7 @@ def marginalize(self, left=0, right=0) -> "Gaussian": P_aa = self.precision[..., a, a] P_ba = self.precision[..., b, a] P_bb = self.precision[..., b, b] - P_b = cholesky(P_bb) + P_b = cholesky(P_bb, jitter=jitter) P_a = triangular_solve(P_ba, P_b, upper=False) P_at = P_a.transpose(-1, -2) precision = P_aa - matmul(P_at, P_a) @@ -403,7 +416,9 @@ def event_permute(self, perm): def __add__(self, other): return self.to_gaussian() + other - def marginalize(self, left=0, right=0): + def marginalize( + self, left: int = 0, right: int = 0, *, jitter: float = 0.0 + ) -> Gaussian: if left == 0 and right == self.loc.size(-1): n = self.matrix.size(-2) precision = self.scale.new_zeros(self.batch_shape + (n, n)) @@ -411,7 +426,7 @@ def marginalize(self, left=0, right=0): log_normalizer = self.scale.new_zeros(self.batch_shape) return Gaussian(log_normalizer, info_vec, precision) else: - return self.to_gaussian().marginalize(left, right) + return self.to_gaussian().marginalize(left, right, jitter=jitter) def mvn_to_gaussian(mvn): @@ -507,7 +522,9 @@ def matrix_and_mvn_to_gaussian(matrix, mvn): return result -def gaussian_tensordot(x: Gaussian, y: Gaussian, dims: int = 0) -> Gaussian: +def gaussian_tensordot( + x: Gaussian, y: Gaussian, dims: int = 0, *, jitter: float = 0.0 +) -> Gaussian: """ Computes the integral over two gaussians: @@ -519,6 +536,8 @@ def gaussian_tensordot(x: Gaussian, y: Gaussian, dims: int = 0) -> Gaussian: :param x: a Gaussian instance :param y: a Gaussian instance :param dims: number of variables to contract + :param float jitter: Optional constant added to matrix diagonals before + performing Cholesky decompositions, to improve stability. """ assert isinstance(x, Gaussian) assert isinstance(y, Gaussian) @@ -550,7 +569,7 @@ def gaussian_tensordot(x: Gaussian, y: Gaussian, dims: int = 0) -> Gaussian: b = xb + yb # Pbb + Qbb needs to be positive definite, so that we can malginalize out `b` (to have a finite integral) - L = cholesky(Pbb + Qbb) + L = cholesky(Pbb + Qbb, jitter=jitter) LinvB = triangular_solve(B, L, upper=False) LinvBt = LinvB.transpose(-2, -1) Linvb = triangular_solve(b.unsqueeze(-1), L, upper=False) @@ -570,13 +589,17 @@ def gaussian_tensordot(x: Gaussian, y: Gaussian, dims: int = 0) -> Gaussian: return Gaussian(log_normalizer, info_vec, precision) -def sequential_gaussian_tensordot(gaussian: Gaussian) -> Gaussian: +def sequential_gaussian_tensordot( + gaussian: Gaussian, *, jitter: float = 0.0 +) -> Gaussian: """ Integrates a Gaussian ``x`` whose rightmost batch dimension is time, computes:: x[..., 0] @ x[..., 1] @ ... @ x[..., T-1] :param Gaussian gaussian: A batched Gaussian whose rightmost dimension is time. + :param float jitter: Optional constant added to matrix diagonals before + performing Cholesky decompositions, to improve stability. :returns: A Markov product of the Gaussian along its time dimension. :rtype: Gaussian """ @@ -590,7 +613,7 @@ def sequential_gaussian_tensordot(gaussian: Gaussian) -> Gaussian: even_part = gaussian[..., :even_time] x_y = even_part.reshape(batch_shape + (even_time // 2, 2)) x, y = x_y[..., 0], x_y[..., 1] - contracted = gaussian_tensordot(x, y, state_dim) + contracted = gaussian_tensordot(x, y, state_dim, jitter=jitter) if time > even_time: contracted = Gaussian.cat((contracted, gaussian[..., -1:]), dim=-1) gaussian = contracted @@ -602,6 +625,8 @@ def sequential_gaussian_filter_sample( trans: Gaussian, sample_shape: Tuple[int, ...] = (), noise: Optional[torch.Tensor] = None, + *, + jitter: float = 0.0, ) -> torch.Tensor: """ Draws a reparameterized sample from a Markov product of Gaussians via @@ -618,6 +643,8 @@ def sequential_gaussian_filter_sample( to be sampled, and ``state_dim = init.dim()`` is the state dimension. This is useful for computing the mean (pass zeros), varying temperature (pass scaled noise), and antithetic sampling (pass ``cat([z,-z])``). + :param float jitter: Optional constant added to matrix diagonals before + performing Cholesky decompositions, to improve stability. :returns: A reparametrized sample of shape ``sample_shape + batch_shape + (duration, state_dim)``. :rtype: torch.Tensor @@ -653,7 +680,7 @@ def sequential_gaussian_filter_sample( y = y.event_pad(left=state_dim) joint = (x + y).event_permute(perm) tape.append(joint) - contracted = joint.marginalize(left=state_dim) + contracted = joint.marginalize(left=state_dim, jitter=jitter) if time > even_time: contracted = Gaussian.cat((contracted, gaussian[..., -1:]), dim=-1) gaussian = contracted diff --git a/pyro/ops/tensor_utils.py b/pyro/ops/tensor_utils.py index 7efc847d59..0084d5cce5 100644 --- a/pyro/ops/tensor_utils.py +++ b/pyro/ops/tensor_utils.py @@ -393,9 +393,11 @@ def inverse_haar_transform(x): return x -def cholesky(x): +def cholesky(x, *, jitter: float = 0.0): if x.size(-1) == 1: return x.sqrt() + if jitter > 0: + x = x + jitter * torch.eye(x.size(-1), device=x.device, dtype=x.dtype) return torch.linalg.cholesky(x) diff --git a/tests/distributions/test_hmm.py b/tests/distributions/test_hmm.py index 01ea316855..00010f94a7 100644 --- a/tests/distributions/test_hmm.py +++ b/tests/distributions/test_hmm.py @@ -433,7 +433,8 @@ def test_gaussian_hmm_distribution( scale = obs_dist.scale_tril.diagonal(dim1=-2, dim2=-1) obs_dist = dist.Normal(obs_dist.loc, scale).to_event(1) d = dist.GaussianHMM( - init_dist, trans_mat, trans_dist, obs_mat, obs_dist, duration=num_steps + init_dist, trans_mat, trans_dist, obs_mat, obs_dist, duration=num_steps, + jitter=1e-8 ) if diag: obs_mvn = dist.MultivariateNormal( diff --git a/tests/ops/test_gaussian.py b/tests/ops/test_gaussian.py index f8fab903b7..30f4717c5a 100644 --- a/tests/ops/test_gaussian.py +++ b/tests/ops/test_gaussian.py @@ -15,6 +15,7 @@ AffineNormal, Gaussian, gaussian_tensordot, + matrix_and_gaussian_to_gaussian, matrix_and_mvn_to_gaussian, mvn_to_gaussian, sequential_gaussian_filter_sample, @@ -378,7 +379,7 @@ def test_gaussian_tensordot( nc = y_dim - dot_dims try: torch.linalg.cholesky(x.precision[..., na:, na:] + y.precision[..., :nb, :nb]) - except RuntimeError: + except Exception: pytest.skip("Cannot marginalize the common variables of two Gaussians.") z = gaussian_tensordot(x, y, dot_dims) @@ -557,3 +558,60 @@ def test_sequential_gaussian_filter_sample_antithetic( ) expected = torch.stack([sample, mean, 2 * mean - sample]) assert torch.allclose(sample3, expected) + + +@pytest.mark.parametrize("num_steps", [10, 100, 1000, 10000, 100000, 1000000]) +def test_sequential_gaussian_filter_sample_stability(num_steps): + # This tests long-chain filtering at low precision. + zero = torch.zeros((), dtype=torch.float) + eye = torch.eye(4, dtype=torch.float) + noise = torch.randn(num_steps, 4, dtype=torch.float, requires_grad=True) + trans_matrix = torch.tensor( + [ + [ + 0.8571434617042542, + -0.23285813629627228, + 0.05360094830393791, + -0.017088839784264565, + ], + [ + 0.7609677314758301, + 0.6596274971961975, + -0.022656921297311783, + 0.05166701227426529, + ], + [ + 3.0979342460632324, + 5.446939945220947, + -0.3425334692001343, + 0.01096670888364315, + ], + [ + -1.8180007934570312, + -0.4965082108974457, + -0.006048532668501139, + -0.08525419235229492, + ], + ], + dtype=torch.float, + requires_grad=True, + ) + + init = Gaussian(zero, zero.expand(4), eye) + trans = matrix_and_gaussian_to_gaussian( + trans_matrix, Gaussian(zero, zero.expand(4), eye) + ).expand((num_steps - 1,)) + + # Check numerically stabilized value. + jitter = 1e-7 + x = sequential_gaussian_filter_sample(init, trans, (), noise, jitter=jitter) + assert torch.isfinite(x).all() + + # Check gradients. + grads = torch.autograd.grad(x.sum(), [trans_matrix, noise]) + assert all(torch.isfinite(g).all() for g in grads) + + if num_steps <= 100: + # Check agreement with unstablized computation. + x_nojitter = sequential_gaussian_filter_sample(init, trans, (), noise) + assert_close(x, x_nojitter, rtol=1e-4, atol=1e-2) From 102abf5cc3db19618e649e7e5cd727339341e884 Mon Sep 17 00:00:00 2001 From: Fritz Obermeyer Date: Thu, 27 Oct 2022 17:00:57 -0700 Subject: [PATCH 02/11] Isolate changes to pyro.ops.tensor_utils --- pyro/distributions/hmm.py | 43 ++++++++--------------------- pyro/ops/gaussian.py | 49 ++++++++------------------------- pyro/ops/tensor_utils.py | 47 +++++++++++++++++++++++++++++-- tests/distributions/test_hmm.py | 3 +- tests/ops/test_gaussian.py | 13 +++++---- 5 files changed, 76 insertions(+), 79 deletions(-) diff --git a/pyro/distributions/hmm.py b/pyro/distributions/hmm.py index 50162e4d29..2a38b19f11 100644 --- a/pyro/distributions/hmm.py +++ b/pyro/distributions/hmm.py @@ -489,8 +489,6 @@ class GaussianHMM(HiddenMarkovModel): :param int duration: Optional size of the time axis ``event_shape[0]``. This is required when sampling from homogeneous HMMs whose parameters are not expanded along the time axis. - :param float jitter: Optional constant added to matrix diagonals before - performing Cholesky decompositions, to improve stability. """ has_rsample = True @@ -506,8 +504,6 @@ def __init__( observation_dist, validate_args=None, duration=None, - *, - jitter: float = 0.0, ): assert isinstance(initial_dist, torch.distributions.MultivariateNormal) or ( isinstance(initial_dist, torch.distributions.Independent) @@ -546,13 +542,11 @@ def __init__( self._init = mvn_to_gaussian(initial_dist).expand(self.batch_shape) self._trans = matrix_and_mvn_to_gaussian(transition_matrix, transition_dist) self._obs = matrix_and_mvn_to_gaussian(observation_matrix, observation_dist) - self.jitter = jitter def expand(self, batch_shape, _instance=None): new = self._get_checked_instance(GaussianHMM, _instance) new.hidden_dim = self.hidden_dim new.obs_dim = self.obs_dim - new.jitter = self.jitter new._obs = self._obs new._trans = self._trans @@ -578,14 +572,10 @@ def log_prob(self, value): ) # Eliminate time dimension. - result = sequential_gaussian_tensordot( - result.expand(result.batch_shape), jitter=self.jitter - ) + result = sequential_gaussian_tensordot(result.expand(result.batch_shape)) # Combine initial factor. - result = gaussian_tensordot( - self._init, result, dims=self.hidden_dim, jitter=self.jitter - ) + result = gaussian_tensordot(self._init, result, dims=self.hidden_dim) # Marginalize out final state. result = result.event_logsumexp() @@ -594,13 +584,11 @@ def log_prob(self, value): def rsample(self, sample_shape=torch.Size()): assert self.duration is not None sample_shape = torch.Size(sample_shape) - trans = self._trans + self._obs.marginalize( - right=self.obs_dim, jitter=self.jitter - ).event_pad(left=self.hidden_dim) - trans = trans.expand(trans.batch_shape[:-1] + (self.duration,)) - z = sequential_gaussian_filter_sample( - self._init, trans, sample_shape, jitter=self.jitter + trans = self._trans + self._obs.marginalize(right=self.obs_dim).event_pad( + left=self.hidden_dim ) + trans = trans.expand(trans.batch_shape[:-1] + (self.duration,)) + z = sequential_gaussian_filter_sample(self._init, trans, sample_shape) z = z[..., 1:, :] # drop the initial hidden state x = self._obs.left_condition(z).rsample() return x @@ -611,9 +599,7 @@ def rsample_posterior(self, value, sample_shape=torch.Size()): """ trans = self._trans + self._obs.condition(value).event_pad(left=self.hidden_dim) trans = trans.expand(trans.batch_shape) - z = sequential_gaussian_filter_sample( - self._init, trans, sample_shape, jitter=self.jitter - ) + z = sequential_gaussian_filter_sample(self._init, trans, sample_shape) z = z[..., 1:, :] # drop the initial hidden state return z @@ -635,14 +621,10 @@ def filter(self, value): logp = self._trans + self._obs.condition(value).event_pad(left=self.hidden_dim) # Eliminate time dimension. - logp = sequential_gaussian_tensordot( - logp.expand(logp.batch_shape), jitter=self.jitter - ) + logp = sequential_gaussian_tensordot(logp.expand(logp.batch_shape)) # Combine initial factor. - logp = gaussian_tensordot( - self._init, logp, dims=self.hidden_dim, jitter=self.jitter - ) + logp = gaussian_tensordot(self._init, logp, dims=self.hidden_dim) # Convert to a distribution precision = logp.precision @@ -682,7 +664,6 @@ def conjugate_update(self, other): new = self._get_checked_instance(GaussianHMM) new.hidden_dim = self.hidden_dim new.obs_dim = self.obs_dim - new.jitter = self.jitter new._init = self._init new._trans = self._trans new._obs = self._obs + mvn_to_gaussian(other.to_event(-1)).event_pad( @@ -691,9 +672,9 @@ def conjugate_update(self, other): # Normalize. # TODO cache this computation for the forward pass of .rsample(). - logp = new._trans + new._obs.marginalize( - right=new.obs_dim, jitter=self.jitter - ).event_pad(left=new.hidden_dim) + logp = new._trans + new._obs.marginalize(right=new.obs_dim).event_pad( + left=new.hidden_dim + ) logp = sequential_gaussian_tensordot(logp.expand(logp.batch_shape)) logp = gaussian_tensordot(new._init, logp, dims=new.hidden_dim) log_normalizer = logp.event_logsumexp() diff --git a/pyro/ops/gaussian.py b/pyro/ops/gaussian.py index fba0ec7a1f..4541513495 100644 --- a/pyro/ops/gaussian.py +++ b/pyro/ops/gaussian.py @@ -9,13 +9,7 @@ from torch.nn.functional import pad from pyro.distributions.util import broadcast_shape -from pyro.ops.tensor_utils import ( - cholesky, - cholesky_solve, - matmul, - matvecmul, - triangular_solve, -) +from pyro.ops.tensor_utils import cholesky, matmul, matvecmul, triangular_solve class Gaussian: @@ -161,7 +155,7 @@ def rsample( Reparameterized sampler. """ P_chol = cholesky(self.precision) - loc = cholesky_solve(self.info_vec.unsqueeze(-1), P_chol).squeeze(-1) + loc = self.info_vec.unsqueeze(-1).cholesky_solve(P_chol).squeeze(-1) shape = sample_shape + self.batch_shape + (self.dim(), 1) if noise is None: noise = torch.randn(shape, dtype=loc.dtype, device=loc.device) @@ -236,9 +230,7 @@ def left_condition(self, value: torch.Tensor) -> "Gaussian": ) return self.event_permute(perm).condition(value) - def marginalize( - self, left: int = 0, right: int = 0, *, jitter: float = 0.0 - ) -> "Gaussian": + def marginalize(self, left=0, right=0) -> "Gaussian": """ Marginalizing out variables on either side of the event dimension:: @@ -249,11 +241,6 @@ def marginalize( g.condition(x).event_logsumexp() = g.marginalize(left=g.dim() - x.size(-1)).log_density(x) - - :param int left: Number of left event dims to marginalize. - :param int right: Number of right event dims to marginalize. - :param float jitter: Optional constant added to matrix diagonals before - performing Cholesky decompositions, to improve stability. """ if left == 0 and right == 0: return self @@ -267,7 +254,7 @@ def marginalize( P_aa = self.precision[..., a, a] P_ba = self.precision[..., b, a] P_bb = self.precision[..., b, b] - P_b = cholesky(P_bb, jitter=jitter) + P_b = cholesky(P_bb) P_a = triangular_solve(P_ba, P_b, upper=False) P_at = P_a.transpose(-1, -2) precision = P_aa - matmul(P_at, P_a) @@ -416,9 +403,7 @@ def event_permute(self, perm): def __add__(self, other): return self.to_gaussian() + other - def marginalize( - self, left: int = 0, right: int = 0, *, jitter: float = 0.0 - ) -> Gaussian: + def marginalize(self, left=0, right=0): if left == 0 and right == self.loc.size(-1): n = self.matrix.size(-2) precision = self.scale.new_zeros(self.batch_shape + (n, n)) @@ -426,7 +411,7 @@ def marginalize( log_normalizer = self.scale.new_zeros(self.batch_shape) return Gaussian(log_normalizer, info_vec, precision) else: - return self.to_gaussian().marginalize(left, right, jitter=jitter) + return self.to_gaussian().marginalize(left, right) def mvn_to_gaussian(mvn): @@ -522,9 +507,7 @@ def matrix_and_mvn_to_gaussian(matrix, mvn): return result -def gaussian_tensordot( - x: Gaussian, y: Gaussian, dims: int = 0, *, jitter: float = 0.0 -) -> Gaussian: +def gaussian_tensordot(x: Gaussian, y: Gaussian, dims: int = 0) -> Gaussian: """ Computes the integral over two gaussians: @@ -536,8 +519,6 @@ def gaussian_tensordot( :param x: a Gaussian instance :param y: a Gaussian instance :param dims: number of variables to contract - :param float jitter: Optional constant added to matrix diagonals before - performing Cholesky decompositions, to improve stability. """ assert isinstance(x, Gaussian) assert isinstance(y, Gaussian) @@ -569,7 +550,7 @@ def gaussian_tensordot( b = xb + yb # Pbb + Qbb needs to be positive definite, so that we can malginalize out `b` (to have a finite integral) - L = cholesky(Pbb + Qbb, jitter=jitter) + L = cholesky(Pbb + Qbb) LinvB = triangular_solve(B, L, upper=False) LinvBt = LinvB.transpose(-2, -1) Linvb = triangular_solve(b.unsqueeze(-1), L, upper=False) @@ -589,17 +570,13 @@ def gaussian_tensordot( return Gaussian(log_normalizer, info_vec, precision) -def sequential_gaussian_tensordot( - gaussian: Gaussian, *, jitter: float = 0.0 -) -> Gaussian: +def sequential_gaussian_tensordot(gaussian: Gaussian) -> Gaussian: """ Integrates a Gaussian ``x`` whose rightmost batch dimension is time, computes:: x[..., 0] @ x[..., 1] @ ... @ x[..., T-1] :param Gaussian gaussian: A batched Gaussian whose rightmost dimension is time. - :param float jitter: Optional constant added to matrix diagonals before - performing Cholesky decompositions, to improve stability. :returns: A Markov product of the Gaussian along its time dimension. :rtype: Gaussian """ @@ -613,7 +590,7 @@ def sequential_gaussian_tensordot( even_part = gaussian[..., :even_time] x_y = even_part.reshape(batch_shape + (even_time // 2, 2)) x, y = x_y[..., 0], x_y[..., 1] - contracted = gaussian_tensordot(x, y, state_dim, jitter=jitter) + contracted = gaussian_tensordot(x, y, state_dim) if time > even_time: contracted = Gaussian.cat((contracted, gaussian[..., -1:]), dim=-1) gaussian = contracted @@ -625,8 +602,6 @@ def sequential_gaussian_filter_sample( trans: Gaussian, sample_shape: Tuple[int, ...] = (), noise: Optional[torch.Tensor] = None, - *, - jitter: float = 0.0, ) -> torch.Tensor: """ Draws a reparameterized sample from a Markov product of Gaussians via @@ -643,8 +618,6 @@ def sequential_gaussian_filter_sample( to be sampled, and ``state_dim = init.dim()`` is the state dimension. This is useful for computing the mean (pass zeros), varying temperature (pass scaled noise), and antithetic sampling (pass ``cat([z,-z])``). - :param float jitter: Optional constant added to matrix diagonals before - performing Cholesky decompositions, to improve stability. :returns: A reparametrized sample of shape ``sample_shape + batch_shape + (duration, state_dim)``. :rtype: torch.Tensor @@ -680,7 +653,7 @@ def sequential_gaussian_filter_sample( y = y.event_pad(left=state_dim) joint = (x + y).event_permute(perm) tape.append(joint) - contracted = joint.marginalize(left=state_dim, jitter=jitter) + contracted = joint.marginalize(left=state_dim) if time > even_time: contracted = Gaussian.cat((contracted, gaussian[..., -1:]), dim=-1) gaussian = contracted diff --git a/pyro/ops/tensor_utils.py b/pyro/ops/tensor_utils.py index 0084d5cce5..39e060f531 100644 --- a/pyro/ops/tensor_utils.py +++ b/pyro/ops/tensor_utils.py @@ -2,11 +2,33 @@ # SPDX-License-Identifier: Apache-2.0 import math +from contextlib import contextmanager +from typing import Optional import torch from torch.fft import irfft, rfft _ROOT_TWO_INVERSE = 1.0 / math.sqrt(2.0) +JITTER = 1e-12 + + +@contextmanager +def settings(*, jitter: float): + """ + Context manager to set global settings. + + :param float jitter: Constant added to matrix diagonals before + performing Cholesky decompositions, to improve stability. + """ + global JITTER + jitter = float(jitter) + assert jitter >= 0 + old = JITTER + try: + JITTER = jitter + yield + finally: + JITTER = old def as_complex(x): @@ -393,11 +415,32 @@ def inverse_haar_transform(x): return x -def cholesky(x, *, jitter: float = 0.0): +def cholesky(x, *, jitter: Optional[float] = None): + if jitter is None: + jitter = JITTER + + # Handle simple scalar case. if x.size(-1) == 1: + if jitter: + x = x.clamp(min=0) return x.sqrt() + if jitter > 0: - x = x + jitter * torch.eye(x.size(-1), device=x.device, dtype=x.dtype) + # First try without jitter. + result, info = torch.linalg.cholesky_ex(x) + if not info.any(): + return result + + # Try adding increasing amounts of jitter where needed. + eye = torch.eye(x.size(-1), device=x.device, dtype=x.dtype) + x = x.clone() + while jitter < 1: + x[info > 0] += jitter * eye + result, info = torch.linalg.cholesky_ex(x) + if not info.any(): + return result + jitter *= 3 + return torch.linalg.cholesky(x) diff --git a/tests/distributions/test_hmm.py b/tests/distributions/test_hmm.py index 00010f94a7..01ea316855 100644 --- a/tests/distributions/test_hmm.py +++ b/tests/distributions/test_hmm.py @@ -433,8 +433,7 @@ def test_gaussian_hmm_distribution( scale = obs_dist.scale_tril.diagonal(dim1=-2, dim2=-1) obs_dist = dist.Normal(obs_dist.loc, scale).to_event(1) d = dist.GaussianHMM( - init_dist, trans_mat, trans_dist, obs_mat, obs_dist, duration=num_steps, - jitter=1e-8 + init_dist, trans_mat, trans_dist, obs_mat, obs_dist, duration=num_steps ) if diag: obs_mvn = dist.MultivariateNormal( diff --git a/tests/ops/test_gaussian.py b/tests/ops/test_gaussian.py index 30f4717c5a..92beec5b14 100644 --- a/tests/ops/test_gaussian.py +++ b/tests/ops/test_gaussian.py @@ -10,6 +10,7 @@ from torch.nn.functional import pad import pyro.distributions as dist +import pyro.ops.tensor_utils from pyro.distributions.util import broadcast_shape from pyro.ops.gaussian import ( AffineNormal, @@ -603,15 +604,15 @@ def test_sequential_gaussian_filter_sample_stability(num_steps): ).expand((num_steps - 1,)) # Check numerically stabilized value. - jitter = 1e-7 - x = sequential_gaussian_filter_sample(init, trans, (), noise, jitter=jitter) + x = sequential_gaussian_filter_sample(init, trans, (), noise) assert torch.isfinite(x).all() # Check gradients. grads = torch.autograd.grad(x.sum(), [trans_matrix, noise]) assert all(torch.isfinite(g).all() for g in grads) - if num_steps <= 100: - # Check agreement with unstablized computation. - x_nojitter = sequential_gaussian_filter_sample(init, trans, (), noise) - assert_close(x, x_nojitter, rtol=1e-4, atol=1e-2) + if num_steps <= 1000: + # Check jitter did not significantly affect computation. + with pyro.ops.tensor_utils.settings(jitter=0.0): + x_nojitter = sequential_gaussian_filter_sample(init, trans, (), noise) + assert_close(x, x_nojitter, rtol=1e-4, atol=1e-2) From ae184e8f38617126effd4d61ed81fb8a534dee89 Mon Sep 17 00:00:00 2001 From: Fritz Obermeyer Date: Thu, 27 Oct 2022 17:19:26 -0700 Subject: [PATCH 03/11] Add a warning --- pyro/ops/tensor_utils.py | 2 ++ tests/ops/test_gaussian.py | 1 + 2 files changed, 3 insertions(+) diff --git a/pyro/ops/tensor_utils.py b/pyro/ops/tensor_utils.py index 39e060f531..5dd7cd9e00 100644 --- a/pyro/ops/tensor_utils.py +++ b/pyro/ops/tensor_utils.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import math +import warnings from contextlib import contextmanager from typing import Optional @@ -432,6 +433,7 @@ def cholesky(x, *, jitter: Optional[float] = None): return result # Try adding increasing amounts of jitter where needed. + warnings.warn("Singular matrix in cholesky(); adding jitter.") eye = torch.eye(x.size(-1), device=x.device, dtype=x.dtype) x = x.clone() while jitter < 1: diff --git a/tests/ops/test_gaussian.py b/tests/ops/test_gaussian.py index 92beec5b14..33dea3506d 100644 --- a/tests/ops/test_gaussian.py +++ b/tests/ops/test_gaussian.py @@ -561,6 +561,7 @@ def test_sequential_gaussian_filter_sample_antithetic( assert torch.allclose(sample3, expected) +@pytest.mark.filterwarnings("ignore:Singular matrix in cholesky") @pytest.mark.parametrize("num_steps", [10, 100, 1000, 10000, 100000, 1000000]) def test_sequential_gaussian_filter_sample_stability(num_steps): # This tests long-chain filtering at low precision. From fd40b54d4d0e3378479847e790bed3add17f21e5 Mon Sep 17 00:00:00 2001 From: Fritz Obermeyer Date: Thu, 27 Oct 2022 17:35:34 -0700 Subject: [PATCH 04/11] Simplify --- pyro/ops/tensor_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyro/ops/tensor_utils.py b/pyro/ops/tensor_utils.py index 5dd7cd9e00..ac203d7db2 100644 --- a/pyro/ops/tensor_utils.py +++ b/pyro/ops/tensor_utils.py @@ -436,8 +436,9 @@ def cholesky(x, *, jitter: Optional[float] = None): warnings.warn("Singular matrix in cholesky(); adding jitter.") eye = torch.eye(x.size(-1), device=x.device, dtype=x.dtype) x = x.clone() + x_diag = x.diagonal(dim1=-1, dim2=-2) while jitter < 1: - x[info > 0] += jitter * eye + x_diag[info > 0] += jitter result, info = torch.linalg.cholesky_ex(x) if not info.any(): return result From d78b83c1057287411d494d2460dda2437b67a14b Mon Sep 17 00:00:00 2001 From: Fritz Obermeyer Date: Thu, 27 Oct 2022 17:35:45 -0700 Subject: [PATCH 05/11] Simplify --- pyro/ops/tensor_utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/pyro/ops/tensor_utils.py b/pyro/ops/tensor_utils.py index ac203d7db2..e88064c6f5 100644 --- a/pyro/ops/tensor_utils.py +++ b/pyro/ops/tensor_utils.py @@ -434,7 +434,6 @@ def cholesky(x, *, jitter: Optional[float] = None): # Try adding increasing amounts of jitter where needed. warnings.warn("Singular matrix in cholesky(); adding jitter.") - eye = torch.eye(x.size(-1), device=x.device, dtype=x.dtype) x = x.clone() x_diag = x.diagonal(dim1=-1, dim2=-2) while jitter < 1: From d542b83fd2130c0c52bf6673bedbc95b62fec9da Mon Sep 17 00:00:00 2001 From: Fritz Obermeyer Date: Thu, 27 Oct 2022 23:12:35 -0700 Subject: [PATCH 06/11] Simplify to fixed jitter based on finfo.eps --- pyro/ops/tensor_utils.py | 52 +++++--------------------------------- tests/ops/test_gaussian.py | 7 ----- 2 files changed, 7 insertions(+), 52 deletions(-) diff --git a/pyro/ops/tensor_utils.py b/pyro/ops/tensor_utils.py index e88064c6f5..eaf77298ca 100644 --- a/pyro/ops/tensor_utils.py +++ b/pyro/ops/tensor_utils.py @@ -2,34 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 import math -import warnings -from contextlib import contextmanager -from typing import Optional import torch from torch.fft import irfft, rfft _ROOT_TWO_INVERSE = 1.0 / math.sqrt(2.0) -JITTER = 1e-12 - - -@contextmanager -def settings(*, jitter: float): - """ - Context manager to set global settings. - - :param float jitter: Constant added to matrix diagonals before - performing Cholesky decompositions, to improve stability. - """ - global JITTER - jitter = float(jitter) - assert jitter >= 0 - old = JITTER - try: - JITTER = jitter - yield - finally: - JITTER = old def as_complex(x): @@ -416,32 +393,17 @@ def inverse_haar_transform(x): return x -def cholesky(x, *, jitter: Optional[float] = None): - if jitter is None: - jitter = JITTER - - # Handle simple scalar case. +def cholesky(x, *, safe: bool = True): if x.size(-1) == 1: - if jitter: - x = x.clamp(min=0) + if safe: + x = x.clamp(min=torch.finfo(x.dtype).tiny) return x.sqrt() - if jitter > 0: - # First try without jitter. - result, info = torch.linalg.cholesky_ex(x) - if not info.any(): - return result - - # Try adding increasing amounts of jitter where needed. - warnings.warn("Singular matrix in cholesky(); adding jitter.") + if safe: x = x.clone() - x_diag = x.diagonal(dim1=-1, dim2=-2) - while jitter < 1: - x_diag[info > 0] += jitter - result, info = torch.linalg.cholesky_ex(x) - if not info.any(): - return result - jitter *= 3 + x_max = x.data.reshape(*x.shape[:-2], -1).abs().max(-1, True).values + jitter = x_max * torch.finfo(x.dtype).eps + x.data.diagonal(dim1=-1, dim2=-2).add_(jitter) return torch.linalg.cholesky(x) diff --git a/tests/ops/test_gaussian.py b/tests/ops/test_gaussian.py index 33dea3506d..0982d8eb91 100644 --- a/tests/ops/test_gaussian.py +++ b/tests/ops/test_gaussian.py @@ -10,7 +10,6 @@ from torch.nn.functional import pad import pyro.distributions as dist -import pyro.ops.tensor_utils from pyro.distributions.util import broadcast_shape from pyro.ops.gaussian import ( AffineNormal, @@ -611,9 +610,3 @@ def test_sequential_gaussian_filter_sample_stability(num_steps): # Check gradients. grads = torch.autograd.grad(x.sum(), [trans_matrix, noise]) assert all(torch.isfinite(g).all() for g in grads) - - if num_steps <= 1000: - # Check jitter did not significantly affect computation. - with pyro.ops.tensor_utils.settings(jitter=0.0): - x_nojitter = sequential_gaussian_filter_sample(init, trans, (), noise) - assert_close(x, x_nojitter, rtol=1e-4, atol=1e-2) From 2f5e18495506568d6f222ef1e0209b2241104e35 Mon Sep 17 00:00:00 2001 From: Fritz Obermeyer Date: Thu, 27 Oct 2022 23:20:13 -0700 Subject: [PATCH 07/11] Simplify --- pyro/ops/tensor_utils.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/pyro/ops/tensor_utils.py b/pyro/ops/tensor_utils.py index eaf77298ca..7ccac7d17a 100644 --- a/pyro/ops/tensor_utils.py +++ b/pyro/ops/tensor_utils.py @@ -393,17 +393,16 @@ def inverse_haar_transform(x): return x -def cholesky(x, *, safe: bool = True): +def cholesky(x): if x.size(-1) == 1: - if safe: - x = x.clamp(min=torch.finfo(x.dtype).tiny) + x = x.clamp(min=torch.finfo(x.dtype).tiny) return x.sqrt() - if safe: - x = x.clone() - x_max = x.data.reshape(*x.shape[:-2], -1).abs().max(-1, True).values - jitter = x_max * torch.finfo(x.dtype).eps - x.data.diagonal(dim1=-1, dim2=-2).add_(jitter) + # Add adaptive jitter. + x = x.clone() + x_max = x.data.reshape(*x.shape[:-2], -1).abs().max(-1, True).values + jitter = x_max * torch.finfo(x.dtype).eps + x.data.diagonal(dim1=-1, dim2=-2).add_(jitter) return torch.linalg.cholesky(x) From fdb1be1ec78cc99da72ebda9cf4862a430b35f2c Mon Sep 17 00:00:00 2001 From: Fritz Obermeyer Date: Fri, 28 Oct 2022 08:01:39 -0700 Subject: [PATCH 08/11] Rename cholesky -> safe_cholesky --- pyro/distributions/hmm.py | 10 +++++----- pyro/distributions/transforms/cholesky.py | 4 ++-- pyro/ops/gaussian.py | 10 +++++----- pyro/ops/tensor_utils.py | 5 +++-- 4 files changed, 15 insertions(+), 14 deletions(-) diff --git a/pyro/distributions/hmm.py b/pyro/distributions/hmm.py index 2a38b19f11..2112d9979b 100644 --- a/pyro/distributions/hmm.py +++ b/pyro/distributions/hmm.py @@ -20,7 +20,7 @@ ) from pyro.ops.indexing import Vindex from pyro.ops.special import safe_log -from pyro.ops.tensor_utils import cholesky, cholesky_solve +from pyro.ops.tensor_utils import cholesky_solve, safe_cholesky from . import constraints from .torch import Categorical, Gamma, Independent, MultivariateNormal @@ -628,9 +628,9 @@ def filter(self, value): # Convert to a distribution precision = logp.precision - loc = cholesky_solve(logp.info_vec.unsqueeze(-1), cholesky(precision)).squeeze( - -1 - ) + loc = cholesky_solve( + logp.info_vec.unsqueeze(-1), safe_cholesky(precision) + ).squeeze(-1) return MultivariateNormal( loc, precision_matrix=precision, validate_args=self._validate_args ) @@ -928,7 +928,7 @@ def filter(self, value): gamma_dist.concentration, gamma_dist.rate, validate_args=self._validate_args ) # Conditional of last state on unit scale - scale_tril = cholesky(logp.precision) + scale_tril = safe_cholesky(logp.precision) loc = cholesky_solve(logp.info_vec.unsqueeze(-1), scale_tril).squeeze(-1) mvn = MultivariateNormal( loc, scale_tril=scale_tril, validate_args=self._validate_args diff --git a/pyro/distributions/transforms/cholesky.py b/pyro/distributions/transforms/cholesky.py index 3b890f5a3b..d2e4d22684 100644 --- a/pyro/distributions/transforms/cholesky.py +++ b/pyro/distributions/transforms/cholesky.py @@ -89,7 +89,7 @@ def log_abs_det_jacobian(self, x, y): class CholeskyTransform(Transform): r""" - Transform via the mapping :math:`y = cholesky(x)`, where `x` is a + Transform via the mapping :math:`y = safe_cholesky(x)`, where `x` is a positive definite matrix. """ bijective = True @@ -116,7 +116,7 @@ def log_abs_det_jacobian(self, x, y): class CorrMatrixCholeskyTransform(CholeskyTransform): r""" - Transform via the mapping :math:`y = cholesky(x)`, where `x` is a + Transform via the mapping :math:`y = safe_cholesky(x)`, where `x` is a correlation matrix. """ bijective = True diff --git a/pyro/ops/gaussian.py b/pyro/ops/gaussian.py index 4541513495..12f17e973c 100644 --- a/pyro/ops/gaussian.py +++ b/pyro/ops/gaussian.py @@ -9,7 +9,7 @@ from torch.nn.functional import pad from pyro.distributions.util import broadcast_shape -from pyro.ops.tensor_utils import cholesky, matmul, matvecmul, triangular_solve +from pyro.ops.tensor_utils import matmul, matvecmul, safe_cholesky, triangular_solve class Gaussian: @@ -154,7 +154,7 @@ def rsample( """ Reparameterized sampler. """ - P_chol = cholesky(self.precision) + P_chol = safe_cholesky(self.precision) loc = self.info_vec.unsqueeze(-1).cholesky_solve(P_chol).squeeze(-1) shape = sample_shape + self.batch_shape + (self.dim(), 1) if noise is None: @@ -254,7 +254,7 @@ def marginalize(self, left=0, right=0) -> "Gaussian": P_aa = self.precision[..., a, a] P_ba = self.precision[..., b, a] P_bb = self.precision[..., b, b] - P_b = cholesky(P_bb) + P_b = safe_cholesky(P_bb) P_a = triangular_solve(P_ba, P_b, upper=False) P_at = P_a.transpose(-1, -2) precision = P_aa - matmul(P_at, P_a) @@ -277,7 +277,7 @@ def event_logsumexp(self) -> torch.Tensor: Integrates out all latent state (i.e. operating on event dimensions). """ n = self.dim() - chol_P = cholesky(self.precision) + chol_P = safe_cholesky(self.precision) chol_P_u = triangular_solve( self.info_vec.unsqueeze(-1), chol_P, upper=False ).squeeze(-1) @@ -550,7 +550,7 @@ def gaussian_tensordot(x: Gaussian, y: Gaussian, dims: int = 0) -> Gaussian: b = xb + yb # Pbb + Qbb needs to be positive definite, so that we can malginalize out `b` (to have a finite integral) - L = cholesky(Pbb + Qbb) + L = safe_cholesky(Pbb + Qbb) LinvB = triangular_solve(B, L, upper=False) LinvBt = LinvB.transpose(-2, -1) Linvb = triangular_solve(b.unsqueeze(-1), L, upper=False) diff --git a/pyro/ops/tensor_utils.py b/pyro/ops/tensor_utils.py index 7ccac7d17a..8a0ba9a9ce 100644 --- a/pyro/ops/tensor_utils.py +++ b/pyro/ops/tensor_utils.py @@ -7,6 +7,7 @@ from torch.fft import irfft, rfft _ROOT_TWO_INVERSE = 1.0 / math.sqrt(2.0) +CHOLESKY_JITTER = 1.0 def as_complex(x): @@ -393,7 +394,7 @@ def inverse_haar_transform(x): return x -def cholesky(x): +def safe_cholesky(x): if x.size(-1) == 1: x = x.clamp(min=torch.finfo(x.dtype).tiny) return x.sqrt() @@ -401,7 +402,7 @@ def cholesky(x): # Add adaptive jitter. x = x.clone() x_max = x.data.reshape(*x.shape[:-2], -1).abs().max(-1, True).values - jitter = x_max * torch.finfo(x.dtype).eps + jitter = CHOLESKY_JITTER * torch.finfo(x.dtype).eps * x_max x.data.diagonal(dim1=-1, dim2=-2).add_(jitter) return torch.linalg.cholesky(x) From 2da3c7585633153baa42474615491fc62ce91626 Mon Sep 17 00:00:00 2001 From: Fritz Obermeyer Date: Fri, 28 Oct 2022 10:35:24 -0700 Subject: [PATCH 09/11] Allow disabling jitter --- pyro/ops/tensor_utils.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/pyro/ops/tensor_utils.py b/pyro/ops/tensor_utils.py index 8a0ba9a9ce..bbfcff458a 100644 --- a/pyro/ops/tensor_utils.py +++ b/pyro/ops/tensor_utils.py @@ -7,7 +7,7 @@ from torch.fft import irfft, rfft _ROOT_TWO_INVERSE = 1.0 / math.sqrt(2.0) -CHOLESKY_JITTER = 1.0 +CHOLESKY_JITTER = 1.0 # in units of finfo.eps def as_complex(x): @@ -396,14 +396,16 @@ def inverse_haar_transform(x): def safe_cholesky(x): if x.size(-1) == 1: - x = x.clamp(min=torch.finfo(x.dtype).tiny) + if CHOLESKY_JITTER: + x = x.clamp(min=torch.finfo(x.dtype).tiny) return x.sqrt() - # Add adaptive jitter. - x = x.clone() - x_max = x.data.reshape(*x.shape[:-2], -1).abs().max(-1, True).values - jitter = CHOLESKY_JITTER * torch.finfo(x.dtype).eps * x_max - x.data.diagonal(dim1=-1, dim2=-2).add_(jitter) + if CHOLESKY_JITTER: + # Add adaptive jitter. + x = x.clone() + x_max = x.data.reshape(*x.shape[:-2], -1).abs().max(-1, True).values + jitter = CHOLESKY_JITTER * torch.finfo(x.dtype).eps * x_max + x.data.diagonal(dim1=-1, dim2=-2).add_(jitter) return torch.linalg.cholesky(x) From 152b9131ced42d96850f6eec3d96f446b18e44ff Mon Sep 17 00:00:00 2001 From: Fritz Obermeyer Date: Sat, 29 Oct 2022 19:20:44 -0700 Subject: [PATCH 10/11] Address review comment --- pyro/ops/tensor_utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pyro/ops/tensor_utils.py b/pyro/ops/tensor_utils.py index bbfcff458a..524c9ddf51 100644 --- a/pyro/ops/tensor_utils.py +++ b/pyro/ops/tensor_utils.py @@ -7,7 +7,7 @@ from torch.fft import irfft, rfft _ROOT_TWO_INVERSE = 1.0 / math.sqrt(2.0) -CHOLESKY_JITTER = 1.0 # in units of finfo.eps +CHOLESKY_RELATIVE_JITTER = 1.0 # in units of finfo.eps def as_complex(x): @@ -396,15 +396,15 @@ def inverse_haar_transform(x): def safe_cholesky(x): if x.size(-1) == 1: - if CHOLESKY_JITTER: + if CHOLESKY_RELATIVE_JITTER: x = x.clamp(min=torch.finfo(x.dtype).tiny) return x.sqrt() - if CHOLESKY_JITTER: + if CHOLESKY_RELATIVE_JITTER: # Add adaptive jitter. x = x.clone() x_max = x.data.reshape(*x.shape[:-2], -1).abs().max(-1, True).values - jitter = CHOLESKY_JITTER * torch.finfo(x.dtype).eps * x_max + jitter = CHOLESKY_RELATIVE_JITTER * torch.finfo(x.dtype).eps * x_max x.data.diagonal(dim1=-1, dim2=-2).add_(jitter) return torch.linalg.cholesky(x) From 3e273e734cf81dd12e2ef5397daf19cb0902a395 Mon Sep 17 00:00:00 2001 From: Fritz Obermeyer Date: Sat, 29 Oct 2022 21:57:16 -0700 Subject: [PATCH 11/11] Switch to column-wise max, increase jitter --- pyro/ops/tensor_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyro/ops/tensor_utils.py b/pyro/ops/tensor_utils.py index 524c9ddf51..f820aa72cf 100644 --- a/pyro/ops/tensor_utils.py +++ b/pyro/ops/tensor_utils.py @@ -7,7 +7,7 @@ from torch.fft import irfft, rfft _ROOT_TWO_INVERSE = 1.0 / math.sqrt(2.0) -CHOLESKY_RELATIVE_JITTER = 1.0 # in units of finfo.eps +CHOLESKY_RELATIVE_JITTER = 4.0 # in units of finfo.eps def as_complex(x): @@ -403,7 +403,7 @@ def safe_cholesky(x): if CHOLESKY_RELATIVE_JITTER: # Add adaptive jitter. x = x.clone() - x_max = x.data.reshape(*x.shape[:-2], -1).abs().max(-1, True).values + x_max = x.data.abs().max(-1).values jitter = CHOLESKY_RELATIVE_JITTER * torch.finfo(x.dtype).eps * x_max x.data.diagonal(dim1=-1, dim2=-2).add_(jitter)