From d88ba872d988b839939c08b4ab53f085fee08c1d Mon Sep 17 00:00:00 2001 From: Javier Burroni Date: Mon, 18 May 2026 13:35:42 -0400 Subject: [PATCH 1/9] Add finite-gradient regression tests for Gaussian mixtures --- tests/distributions/test_gaussian_mixtures.py | 75 +++++++++++++++++++ 1 file changed, 75 insertions(+) diff --git a/tests/distributions/test_gaussian_mixtures.py b/tests/distributions/test_gaussian_mixtures.py index 7623d4e14b..086b81bbe7 100644 --- a/tests/distributions/test_gaussian_mixtures.py +++ b/tests/distributions/test_gaussian_mixtures.py @@ -209,6 +209,81 @@ def test_gsm_log_prob(): assert_equal( log_prob, correct_log_prob, msg="bad log prob for GaussianScaleMixture" ) + + +def test_mix_of_diag_normals_backward_small_scales_finite_grads(): + torch.manual_seed(0) + + K = 5 + D = 50 + + locs = torch.randn(K, D, requires_grad=True) + coord_scale = torch.full((K, D), 0.1, dtype=torch.float32, requires_grad=True) + component_logits = torch.zeros(K, requires_grad=True) + + dist = MixtureOfDiagNormals(locs, coord_scale, component_logits) + z = dist.rsample(sample_shape=(32,)) + + assert torch.isfinite(z).all() + + loss = z.pow(2).sum() + assert torch.isfinite(loss) + + loss.backward() + + assert torch.isfinite(locs.grad).all() + assert torch.isfinite(coord_scale.grad).all() + assert torch.isfinite(component_logits.grad).all() + +def test_gsm_backward_small_coord_scale_finite_grads(): + torch.manual_seed(0) + + K = 5 + D = 50 + + coord_scale = torch.full((D,), 0.1, requires_grad=True) + component_logits = torch.zeros(K, requires_grad=True) + component_scale = torch.ones(K, dtype=torch.float32, requires_grad=True) + + dist = GaussianScaleMixture(coord_scale, component_logits, component_scale) + z = dist.rsample(sample_shape=(32,)) + + assert torch.isfinite(z).all() + + loss = z.pow(2).sum() + assert torch.isfinite(loss) + + loss.backward() + + assert torch.isfinite(coord_scale.grad).all() + assert torch.isfinite(component_logits.grad).all() + assert torch.isfinite(component_scale.grad).all() + +def test_mix_of_diag_normals_shared_cov_backward_high_dim_finite_grads(): + torch.manual_seed(0) + + K = 3 + D = 300 + + locs = torch.randn(K, D, requires_grad=True) + coord_scale = torch.ones(D, dtype=torch.float32, requires_grad=True) + component_logits = torch.zeros(K, requires_grad=True) + + dist = MixtureOfDiagNormalsSharedCovariance( + locs, coord_scale, component_logits + ) + z = dist.rsample(sample_shape=(4,)) + + assert torch.isfinite(z).all() + + loss = z.pow(2).sum() + assert torch.isfinite(loss) + + loss.backward() + + assert torch.isfinite(locs.grad).all() + assert torch.isfinite(coord_scale.grad).all() + assert torch.isfinite(component_logits.grad).all() @pytest.mark.parametrize("batch_size", [1, 3]) From 980723e0a9cb3bd69b475a21db6f41d7dd93da29 Mon Sep 17 00:00:00 2001 From: Javier Burroni Date: Mon, 18 May 2026 13:40:17 -0400 Subject: [PATCH 2/9] Stabilize Gaussian mixture backward computations --- pyro/distributions/diag_normal_mixture.py | 12 ++-- .../diag_normal_mixture_shared_cov.py | 7 +-- pyro/distributions/gaussian_scale_mixture.py | 58 +++++++++++-------- 3 files changed, 44 insertions(+), 33 deletions(-) diff --git a/pyro/distributions/diag_normal_mixture.py b/pyro/distributions/diag_normal_mixture.py index 52b8725448..c956eeade4 100644 --- a/pyro/distributions/diag_normal_mixture.py +++ b/pyro/distributions/diag_normal_mixture.py @@ -202,23 +202,23 @@ def backward(ctx, grad_output): root_two = math.sqrt(2.0) shift_log_scales = log_scales[..., shift_indices] shift_log_scales[..., 0] = 0.0 - sigma_products = torch.cumsum(shift_log_scales, dim=-1).exp() # b j i + log_sigma_products = torch.cumsum(shift_log_scales, dim=-1) # b j i reverse_indices = torch.tensor( range(dim - 1, -1, -1), dtype=torch.long, device=z.device ) reverse_log_sigma_0 = sigma_0.log()[..., reverse_indices] # b 1 i - sigma_0_products = torch.cumsum(reverse_log_sigma_0, dim=-1).exp()[ + log_sigma_0_products = torch.cumsum(reverse_log_sigma_0, dim=-1)[ ..., reverse_indices - 1 ] # b 1 i - sigma_0_products[..., -1] = 1.0 - sigma_products *= sigma_0_products + log_sigma_0_products[..., -1] = 0.0 + log_sigma_products += log_sigma_0_products logits_grad = torch.erf(z_tilde / root_two) - torch.erf( z_shift / root_two ) # l b j i - logits_grad *= torch.exp(-0.5 * r_sqr_ji) # l b j i - logits_grad = (logits_grad * g / sigma_products).sum(-1) # l b j + logits_grad *= torch.exp(-0.5 * r_sqr_ji - log_sigma_products) # l b j i + logits_grad = (logits_grad * g).sum(-1) # l b j logits_grad = sum_leftmost(logits_grad / q_tot, -1 - batch_dims) # b j logits_grad *= 0.5 * math.pow(2.0 * math.pi, -0.5 * (dim - 1)) logits_grad = -pis * logits_grad diff --git a/pyro/distributions/diag_normal_mixture_shared_cov.py b/pyro/distributions/diag_normal_mixture_shared_cov.py index 5e5b9fa9aa..389df69acc 100644 --- a/pyro/distributions/diag_normal_mixture_shared_cov.py +++ b/pyro/distributions/diag_normal_mixture_shared_cov.py @@ -176,9 +176,7 @@ def backward(ctx, grad_output): log_q_j = log_qs.sum(-1, keepdim=True) # l b j 1 log_q_j_max = torch.max(log_q_j, -2, keepdim=True)[0] q_j_prime = torch.exp(log_q_j - log_q_j_max) # l b j 1 - q_j = torch.exp(log_q_j) # l b j 1 - q_tot = (pis.unsqueeze(-1) * q_j).sum(-2) # l b 1 q_tot_prime = (pis.unsqueeze(-1) * q_j_prime).sum(-2).unsqueeze(-1) # l b 1 1 root_two = math.sqrt(2.0) @@ -186,7 +184,8 @@ def backward(ctx, grad_output): logits_grad = torch.erf((z_ll_ab - mu_ll_ab) / root_two) - torch.erf( (z_ll_ab + mu_ll_ba) / root_two ) - logits_grad *= torch.exp(-0.5 * z_perp_ab_sqr) # l b k j + log_q_j_max = log_q_j_max.squeeze(-1).unsqueeze(-1) # l b 1 1 + logits_grad *= torch.exp(-0.5 * z_perp_ab_sqr - log_q_j_max) # l b k j # bi lbi bkji mu_ab_sigma_g = ((coord_scale * g).unsqueeze(-2).unsqueeze(-2) * mu_ab).sum( @@ -194,7 +193,7 @@ def backward(ctx, grad_output): ) # l b k j logits_grad *= -mu_ab_sigma_g * pis.unsqueeze(-2) # l b k j logits_grad = pis * sum_leftmost( - logits_grad.sum(-1) / q_tot, -(1 + batch_dims) + logits_grad.sum(-1) / q_tot_prime.squeeze(-1), -(1 + batch_dims) ) # b k logits_grad *= math.sqrt(0.5 * math.pi) diff --git a/pyro/distributions/gaussian_scale_mixture.py b/pyro/distributions/gaussian_scale_mixture.py index f1dc47cf64..48e7c77fd6 100644 --- a/pyro/distributions/gaussian_scale_mixture.py +++ b/pyro/distributions/gaussian_scale_mixture.py @@ -151,21 +151,34 @@ def backward(ctx, grad_output): epsilons_sqr = torch.pow(epsilons, 2.0) # l i r_sqr = epsilons_sqr.sum(-1, keepdim=True) # l r_sqr_j = r_sqr / component_scale_sqr # l j - coord_scale_product = coord_scale.prod() - component_scale_power = torch.pow(component_scale, float(dim)) + log_coord_scale_product = coord_scale.log().sum() + log_component_scale_power = component_scale.log() * float(dim) + + log_q_j = ( + -0.5 * r_sqr_j + - 0.5 * math.log(2.0 * math.pi) * float(dim) + - log_coord_scale_product + - log_component_scale_power + ) + log_q_tot = torch.logsumexp(pis.log() + log_q_j, dim=-1, keepdim=True) + q_j_over_q_tot = torch.exp(log_q_j - log_q_tot) + + log_normalizer = ( + -0.5 * math.log(2.0 * math.pi) * float(dim) + - log_component_scale_power + - log_coord_scale_product + - log_q_tot + ) - q_j = torch.exp(-0.5 * r_sqr_j) / math.pow( - 2.0 * math.pi, 0.5 * float(dim) - ) # l j - q_j /= coord_scale_product * component_scale_power # l j - q_tot = (pis * q_j).sum(-1, keepdim=True) # l + Phi_j = torch.exp(-0.5 * r_sqr_j + log_normalizer) - Phi_j = torch.exp(-0.5 * r_sqr_j) # l j - exponents = -torch.arange(1.0, int(dim / 2) + 1.0, 1.0) + exponents = -torch.arange( + 1.0, int(dim / 2) + 1.0, 1.0, device=z.device, dtype=z.dtype + ) if z.dim() > 1: - r_j_poly = r_sqr_j.unsqueeze(-1).expand(-1, -1, int(dim / 2)) # l j d/2 + r_j_poly = r_sqr_j.unsqueeze(-1).expand(-1, -1, int(dim / 2)) else: - r_j_poly = r_sqr_j.unsqueeze(-1).expand(-1, int(dim / 2)) # l j d/2 + r_j_poly = r_sqr_j.unsqueeze(-1).expand(-1, int(dim / 2)) r_j_poly = coeffs * torch.pow(r_j_poly, exponents) Phi_j *= r_j_poly.sum(-1) if dim % 2 == 1: @@ -174,20 +187,19 @@ def backward(ctx, grad_output): coeffs[-1] * math.sqrt(0.5 * math.pi) * (1.0 - torch.erf(r_sqr_j.sqrt() / root_two)) - ) # l j - Phi_j += extra_term * torch.pow(r_sqr_j, -0.5 * float(dim)) - - logits_grad = (z.unsqueeze(-2) * Phi_j.unsqueeze(-1) * g).sum(-1) # l j - logits_grad /= q_tot - logits_grad = sum_leftmost(logits_grad, -1) * math.pow( - 2.0 * math.pi, -0.5 * float(dim) - ) - logits_grad = pis * logits_grad / (component_scale_power * coord_scale_product) + ) + Phi_j += ( + extra_term + * torch.pow(r_sqr_j, -0.5 * float(dim)) + * torch.exp(log_normalizer) + ) + + logits_grad = (z.unsqueeze(-2) * Phi_j.unsqueeze(-1) * g).sum(-1) + logits_grad = sum_leftmost(logits_grad, -1) + logits_grad = pis * logits_grad logits_grad = logits_grad - logits_grad.sum() * pis - prefactor = ( - pis.unsqueeze(-1) * q_j.unsqueeze(-1) * g / q_tot.unsqueeze(-1) - ) # l j i + prefactor = pis.unsqueeze(-1) * q_j_over_q_tot.unsqueeze(-1) * g # l j i coord_scale_grad = sum_leftmost(prefactor * epsilons.unsqueeze(-2), -1) component_scale_grad = sum_leftmost( (prefactor * z.unsqueeze(-2)).sum(-1) / component_scale, -1 From 27f6d2d69c9b9e424452c710ac9827d2e3418f67 Mon Sep 17 00:00:00 2001 From: Javier Burroni Date: Fri, 5 Jun 2026 15:25:08 -0400 Subject: [PATCH 3/9] Format Gaussian mixture code with black --- pyro/distributions/gaussian_scale_mixture.py | 2 +- tests/distributions/test_gaussian_mixtures.py | 14 +++++++------- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/pyro/distributions/gaussian_scale_mixture.py b/pyro/distributions/gaussian_scale_mixture.py index 48e7c77fd6..645bed0b48 100644 --- a/pyro/distributions/gaussian_scale_mixture.py +++ b/pyro/distributions/gaussian_scale_mixture.py @@ -199,7 +199,7 @@ def backward(ctx, grad_output): logits_grad = pis * logits_grad logits_grad = logits_grad - logits_grad.sum() * pis - prefactor = pis.unsqueeze(-1) * q_j_over_q_tot.unsqueeze(-1) * g # l j i + prefactor = pis.unsqueeze(-1) * q_j_over_q_tot.unsqueeze(-1) * g # l j i coord_scale_grad = sum_leftmost(prefactor * epsilons.unsqueeze(-2), -1) component_scale_grad = sum_leftmost( (prefactor * z.unsqueeze(-2)).sum(-1) / component_scale, -1 diff --git a/tests/distributions/test_gaussian_mixtures.py b/tests/distributions/test_gaussian_mixtures.py index 086b81bbe7..32e892441b 100644 --- a/tests/distributions/test_gaussian_mixtures.py +++ b/tests/distributions/test_gaussian_mixtures.py @@ -209,8 +209,8 @@ def test_gsm_log_prob(): assert_equal( log_prob, correct_log_prob, msg="bad log prob for GaussianScaleMixture" ) - - + + def test_mix_of_diag_normals_backward_small_scales_finite_grads(): torch.manual_seed(0) @@ -234,7 +234,8 @@ def test_mix_of_diag_normals_backward_small_scales_finite_grads(): assert torch.isfinite(locs.grad).all() assert torch.isfinite(coord_scale.grad).all() assert torch.isfinite(component_logits.grad).all() - + + def test_gsm_backward_small_coord_scale_finite_grads(): torch.manual_seed(0) @@ -258,7 +259,8 @@ def test_gsm_backward_small_coord_scale_finite_grads(): assert torch.isfinite(coord_scale.grad).all() assert torch.isfinite(component_logits.grad).all() assert torch.isfinite(component_scale.grad).all() - + + def test_mix_of_diag_normals_shared_cov_backward_high_dim_finite_grads(): torch.manual_seed(0) @@ -269,9 +271,7 @@ def test_mix_of_diag_normals_shared_cov_backward_high_dim_finite_grads(): coord_scale = torch.ones(D, dtype=torch.float32, requires_grad=True) component_logits = torch.zeros(K, requires_grad=True) - dist = MixtureOfDiagNormalsSharedCovariance( - locs, coord_scale, component_logits - ) + dist = MixtureOfDiagNormalsSharedCovariance(locs, coord_scale, component_logits) z = dist.rsample(sample_shape=(4,)) assert torch.isfinite(z).all() From 7b8910b91b55b95b1a0ed2dff29f21c74f7b530e Mon Sep 17 00:00:00 2001 From: Javier Burroni Date: Wed, 10 Jun 2026 14:05:55 -0400 Subject: [PATCH 4/9] removed manual_seed --- tests/distributions/test_gaussian_mixtures.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/tests/distributions/test_gaussian_mixtures.py b/tests/distributions/test_gaussian_mixtures.py index 32e892441b..add867ec95 100644 --- a/tests/distributions/test_gaussian_mixtures.py +++ b/tests/distributions/test_gaussian_mixtures.py @@ -212,8 +212,6 @@ def test_gsm_log_prob(): def test_mix_of_diag_normals_backward_small_scales_finite_grads(): - torch.manual_seed(0) - K = 5 D = 50 @@ -237,8 +235,6 @@ def test_mix_of_diag_normals_backward_small_scales_finite_grads(): def test_gsm_backward_small_coord_scale_finite_grads(): - torch.manual_seed(0) - K = 5 D = 50 @@ -262,8 +258,6 @@ def test_gsm_backward_small_coord_scale_finite_grads(): def test_mix_of_diag_normals_shared_cov_backward_high_dim_finite_grads(): - torch.manual_seed(0) - K = 3 D = 300 From 56f1d1f472f9c6b14aa50462d48b1fde987ba314 Mon Sep 17 00:00:00 2001 From: Javier Burroni Date: Wed, 10 Jun 2026 14:59:46 -0400 Subject: [PATCH 5/9] simplify changes --- pyro/distributions/diag_normal_mixture_shared_cov.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/pyro/distributions/diag_normal_mixture_shared_cov.py b/pyro/distributions/diag_normal_mixture_shared_cov.py index 389df69acc..8a4c13ba7d 100644 --- a/pyro/distributions/diag_normal_mixture_shared_cov.py +++ b/pyro/distributions/diag_normal_mixture_shared_cov.py @@ -174,17 +174,16 @@ def backward(ctx, grad_output): epsilons = z_tilde.unsqueeze(-2) - locs_tilde # l b j i log_qs = -0.5 * torch.pow(epsilons, 2.0) # l b j i log_q_j = log_qs.sum(-1, keepdim=True) # l b j 1 - log_q_j_max = torch.max(log_q_j, -2, keepdim=True)[0] + log_q_j_max = torch.max(log_q_j, -2, keepdim=True)[0] # l b 1 1 q_j_prime = torch.exp(log_q_j - log_q_j_max) # l b j 1 - q_tot_prime = (pis.unsqueeze(-1) * q_j_prime).sum(-2).unsqueeze(-1) # l b 1 1 + q_tot_prime = (pis.unsqueeze(-1) * q_j_prime).sum(-2) # l b 1 root_two = math.sqrt(2.0) mu_ll_ba = torch.transpose(mu_ll_ab, -1, -2) logits_grad = torch.erf((z_ll_ab - mu_ll_ab) / root_two) - torch.erf( (z_ll_ab + mu_ll_ba) / root_two ) - log_q_j_max = log_q_j_max.squeeze(-1).unsqueeze(-1) # l b 1 1 logits_grad *= torch.exp(-0.5 * z_perp_ab_sqr - log_q_j_max) # l b k j # bi lbi bkji @@ -193,13 +192,13 @@ def backward(ctx, grad_output): ) # l b k j logits_grad *= -mu_ab_sigma_g * pis.unsqueeze(-2) # l b k j logits_grad = pis * sum_leftmost( - logits_grad.sum(-1) / q_tot_prime.squeeze(-1), -(1 + batch_dims) + logits_grad.sum(-1) / q_tot_prime, -(1 + batch_dims) ) # b k logits_grad *= math.sqrt(0.5 * math.pi) # b j l b j 1 l b i l b 1 1 prefactor = ( - pis.unsqueeze(-1) * q_j_prime * g.unsqueeze(-2) / q_tot_prime + pis.unsqueeze(-1) * q_j_prime * g.unsqueeze(-2) / q_tot_prime.unsqueeze(-1) ) # l b j i locs_grad = sum_leftmost(prefactor, -(2 + batch_dims)) # b j i coord_scale_grad = sum_leftmost(prefactor * epsilons, -(2 + batch_dims)).sum( From 65548957d7b2fe22a70af8c1ba21863ca821fc1d Mon Sep 17 00:00:00 2001 From: Javier Burroni Date: Wed, 10 Jun 2026 16:07:23 -0400 Subject: [PATCH 6/9] simplify changes --- pyro/distributions/gaussian_scale_mixture.py | 53 ++++++++------------ 1 file changed, 22 insertions(+), 31 deletions(-) diff --git a/pyro/distributions/gaussian_scale_mixture.py b/pyro/distributions/gaussian_scale_mixture.py index 645bed0b48..d29a7d6de1 100644 --- a/pyro/distributions/gaussian_scale_mixture.py +++ b/pyro/distributions/gaussian_scale_mixture.py @@ -149,36 +149,29 @@ def backward(ctx, grad_output): component_scale_sqr = torch.pow(component_scale, 2.0) # j epsilons = z / coord_scale # l i epsilons_sqr = torch.pow(epsilons, 2.0) # l i - r_sqr = epsilons_sqr.sum(-1, keepdim=True) # l + r_sqr = epsilons_sqr.sum(-1, keepdim=True) # l 1 r_sqr_j = r_sqr / component_scale_sqr # l j log_coord_scale_product = coord_scale.log().sum() log_component_scale_power = component_scale.log() * float(dim) - - log_q_j = ( - -0.5 * r_sqr_j - - 0.5 * math.log(2.0 * math.pi) * float(dim) - - log_coord_scale_product - - log_component_scale_power - ) - log_q_tot = torch.logsumexp(pis.log() + log_q_j, dim=-1, keepdim=True) - q_j_over_q_tot = torch.exp(log_q_j - log_q_tot) - - log_normalizer = ( - -0.5 * math.log(2.0 * math.pi) * float(dim) - - log_component_scale_power - - log_coord_scale_product - - log_q_tot - ) - - Phi_j = torch.exp(-0.5 * r_sqr_j + log_normalizer) + + log_gaussian_normalizer = ( + 0.5 * math.log(2.0 * math.pi) * float(dim) + + log_coord_scale_product + + log_component_scale_power + ).unsqueeze(0) # 1 j + + log_q_j = -0.5 * r_sqr_j - log_gaussian_normalizer # l j + log_q_tot = torch.logsumexp(pis.log() + log_q_j, dim=-1, keepdim=True) # l 1 + posterior_j = torch.exp(pis.log() + log_q_j - log_q_tot) # l j + Phi_j = torch.exp(-0.5 * r_sqr_j) # l j exponents = -torch.arange( 1.0, int(dim / 2) + 1.0, 1.0, device=z.device, dtype=z.dtype ) if z.dim() > 1: - r_j_poly = r_sqr_j.unsqueeze(-1).expand(-1, -1, int(dim / 2)) + r_j_poly = r_sqr_j.unsqueeze(-1).expand(-1, -1, int(dim / 2)) # l j d/2 else: - r_j_poly = r_sqr_j.unsqueeze(-1).expand(-1, int(dim / 2)) + r_j_poly = r_sqr_j.unsqueeze(-1).expand(-1, int(dim / 2)) # l j d/2 r_j_poly = coeffs * torch.pow(r_j_poly, exponents) Phi_j *= r_j_poly.sum(-1) if dim % 2 == 1: @@ -186,20 +179,18 @@ def backward(ctx, grad_output): extra_term = ( coeffs[-1] * math.sqrt(0.5 * math.pi) - * (1.0 - torch.erf(r_sqr_j.sqrt() / root_two)) - ) - Phi_j += ( - extra_term - * torch.pow(r_sqr_j, -0.5 * float(dim)) - * torch.exp(log_normalizer) - ) - - logits_grad = (z.unsqueeze(-2) * Phi_j.unsqueeze(-1) * g).sum(-1) + * torch.erfc(r_sqr_j.sqrt() / root_two)) # l j + Phi_j += extra_term * torch.pow(r_sqr_j, -0.5 * float(dim)) + + logits_grad = (z.unsqueeze(-2) * Phi_j.unsqueeze(-1) * g).sum(-1) # l j + + log_logits_scale = -log_gaussian_normalizer - log_q_tot # l j + logits_grad = logits_grad * torch.exp(log_logits_scale) # l j logits_grad = sum_leftmost(logits_grad, -1) logits_grad = pis * logits_grad logits_grad = logits_grad - logits_grad.sum() * pis - prefactor = pis.unsqueeze(-1) * q_j_over_q_tot.unsqueeze(-1) * g # l j i + prefactor = posterior_j.unsqueeze(-1) * g # l j i coord_scale_grad = sum_leftmost(prefactor * epsilons.unsqueeze(-2), -1) component_scale_grad = sum_leftmost( (prefactor * z.unsqueeze(-2)).sum(-1) / component_scale, -1 From b6f9f5b0bbbb46228ceaef250439fd50b82f0bd6 Mon Sep 17 00:00:00 2001 From: Javier Burroni Date: Wed, 10 Jun 2026 21:21:29 -0400 Subject: [PATCH 7/9] Format code to pass lint --- pyro/distributions/gaussian_scale_mixture.py | 14 ++++++++------ uv.lock | 3 +++ 2 files changed, 11 insertions(+), 6 deletions(-) create mode 100644 uv.lock diff --git a/pyro/distributions/gaussian_scale_mixture.py b/pyro/distributions/gaussian_scale_mixture.py index d29a7d6de1..3dab79c4ba 100644 --- a/pyro/distributions/gaussian_scale_mixture.py +++ b/pyro/distributions/gaussian_scale_mixture.py @@ -153,14 +153,15 @@ def backward(ctx, grad_output): r_sqr_j = r_sqr / component_scale_sqr # l j log_coord_scale_product = coord_scale.log().sum() log_component_scale_power = component_scale.log() * float(dim) - + log_gaussian_normalizer = ( 0.5 * math.log(2.0 * math.pi) * float(dim) + log_coord_scale_product + log_component_scale_power - ).unsqueeze(0) # 1 j + ) + log_gaussian_normalizer = log_gaussian_normalizer.unsqueeze(0) # 1 j - log_q_j = -0.5 * r_sqr_j - log_gaussian_normalizer # l j + log_q_j = -0.5 * r_sqr_j - log_gaussian_normalizer # l j log_q_tot = torch.logsumexp(pis.log() + log_q_j, dim=-1, keepdim=True) # l 1 posterior_j = torch.exp(pis.log() + log_q_j - log_q_tot) # l j Phi_j = torch.exp(-0.5 * r_sqr_j) # l j @@ -179,12 +180,13 @@ def backward(ctx, grad_output): extra_term = ( coeffs[-1] * math.sqrt(0.5 * math.pi) - * torch.erfc(r_sqr_j.sqrt() / root_two)) # l j + * torch.erfc(r_sqr_j.sqrt() / root_two) + ) # l j Phi_j += extra_term * torch.pow(r_sqr_j, -0.5 * float(dim)) logits_grad = (z.unsqueeze(-2) * Phi_j.unsqueeze(-1) * g).sum(-1) # l j - - log_logits_scale = -log_gaussian_normalizer - log_q_tot # l j + + log_logits_scale = -log_gaussian_normalizer - log_q_tot # l j logits_grad = logits_grad * torch.exp(log_logits_scale) # l j logits_grad = sum_leftmost(logits_grad, -1) logits_grad = pis * logits_grad diff --git a/uv.lock b/uv.lock new file mode 100644 index 0000000000..7518fc90bf --- /dev/null +++ b/uv.lock @@ -0,0 +1,3 @@ +version = 1 +revision = 3 +requires-python = ">=3.12" From 1af3ada61c2d89581ac07bda7d70fb92c971135d Mon Sep 17 00:00:00 2001 From: Javier Burroni Date: Wed, 8 Jul 2026 16:00:28 -0400 Subject: [PATCH 8/9] add a test for the "odd" correction of eq 43 --- tests/distributions/test_gaussian_mixtures.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/distributions/test_gaussian_mixtures.py b/tests/distributions/test_gaussian_mixtures.py index add867ec95..23d5b0e904 100644 --- a/tests/distributions/test_gaussian_mixtures.py +++ b/tests/distributions/test_gaussian_mixtures.py @@ -22,7 +22,7 @@ [MixtureOfDiagNormals, MixtureOfDiagNormalsSharedCovariance, GaussianScaleMixture], ) @pytest.mark.parametrize("K", [3]) -@pytest.mark.parametrize("D", [2, 4]) +@pytest.mark.parametrize("D", [2, 3, 4]) @pytest.mark.parametrize("batch_mode", [True, False]) @pytest.mark.parametrize("flat_logits", [True, False]) @pytest.mark.parametrize("cost_function", ["quadratic"]) From da9d9d44665fca39cd000736d73abf1b71506333 Mon Sep 17 00:00:00 2001 From: Javier Burroni Date: Thu, 9 Jul 2026 15:14:03 -0400 Subject: [PATCH 9/9] GSM: add a comment relating the code to Eq. 43 --- pyro/distributions/gaussian_scale_mixture.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyro/distributions/gaussian_scale_mixture.py b/pyro/distributions/gaussian_scale_mixture.py index 3dab79c4ba..c28c5aa58b 100644 --- a/pyro/distributions/gaussian_scale_mixture.py +++ b/pyro/distributions/gaussian_scale_mixture.py @@ -175,6 +175,8 @@ def backward(ctx, grad_output): r_j_poly = r_sqr_j.unsqueeze(-1).expand(-1, int(dim / 2)) # l j d/2 r_j_poly = coeffs * torch.pow(r_j_poly, exponents) Phi_j *= r_j_poly.sum(-1) + # Eq. 43 odd-D tail term, written with r_sqr_j = r_j**2 + # Power -D/2 appears because Phi_j is later multiplied by z, not r_hat. if dim % 2 == 1: root_two = math.sqrt(2.0) extra_term = (