From bce8505ce265df250f5fc2d3cbb9d086fb294da3 Mon Sep 17 00:00:00 2001 From: scraed Date: Tue, 13 Jan 2026 18:29:06 +0800 Subject: [PATCH] Revert "Avoid MultivariateNormal crash on non-finite dynamics" This reverts commit 9432a34c38631d903ab2b3af0f7941317d70291a. --- src/LanPaint/utils.py | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/src/LanPaint/utils.py b/src/LanPaint/utils.py index ab990af..774ad4c 100644 --- a/src/LanPaint/utils.py +++ b/src/LanPaint/utils.py @@ -264,12 +264,6 @@ class StochasticHarmonicOscillator: cov_yv = (zeta2(Gamma_hat, Delta) * Gamma_hat * D ) **2 / 2 / (Gamma ** 0.5) # sample new position and velocity with multivariate normal distribution - tol = 1e-8 - y_mean = torch.where(torch.isfinite(y_mean), y_mean, y0) - v_mean = torch.where(torch.isfinite(v_mean), v_mean, v0) - cov_yy = torch.where(torch.isfinite(cov_yy), cov_yy, torch.full_like(cov_yy, tol)) - cov_vv = torch.where(torch.isfinite(cov_vv), cov_vv, torch.full_like(cov_vv, tol)) - cov_yv = torch.where(torch.isfinite(cov_yv), cov_yv, torch.zeros_like(cov_yv)) batch_shape = y0.shape cov_matrix = torch.zeros(*batch_shape, 2, 2, device=y0.device, dtype=y0.dtype) @@ -283,6 +277,7 @@ class StochasticHarmonicOscillator: # Compute the Cholesky decomposition to get scale_tril #scale_tril = torch.linalg.cholesky(cov_matrix) scale_tril = torch.zeros(*batch_shape, 2, 2, device=y0.device, dtype=y0.dtype) + tol = 1e-8 cov_yy = torch.clamp( cov_yy, min = tol ) sd_yy = torch.sqrt( cov_yy ) inv_sd_yy = 1/(sd_yy) @@ -291,9 +286,6 @@ class StochasticHarmonicOscillator: scale_tril[..., 0, 1] = 0. scale_tril[..., 1, 0] = cov_yv * inv_sd_yy scale_tril[..., 1, 1] = torch.clamp( cov_vv - cov_yv**2 / cov_yy, min = tol ) ** 0.5 - scale_tril = torch.where(torch.isfinite(scale_tril), scale_tril, torch.zeros_like(scale_tril)) - scale_tril[..., 0, 0] = torch.clamp(scale_tril[..., 0, 0], min=tol) - scale_tril[..., 1, 1] = torch.clamp(scale_tril[..., 1, 1], min=tol) # check if it matches torch.linalg. #assert torch.allclose(torch.linalg.cholesky(cov_matrix), scale_tril, atol = 1e-4, rtol = 1e-4 ) # Sample correlated noise from multivariate normal