Stabilize Gaussian mixture backward passes

[jburroni]

Summary

This PR fixes numerical instabilities in the backward passes for Pyro's Gaussian mixture distributions:

• MixtureOfDiagNormals

• GaussianScaleMixture

• MixtureOfDiagNormalsSharedCovariance

The affected backward paths can materialize very small density/product terms in probability space. In float32, these terms can underflow to zero, which can then produce nonfinite gradients during backpropagation.

Changes

• Adds regression tests that assert finite gradients in numerically sensitive high-dimensional / small-scale cases.

• Updates the affected backward computations to avoid unstable probability-space divisions/products where possible.

• Format code to pass lint.

Tests

Added regression coverage for:

• small-scale MixtureOfDiagNormals

• small-scale GaussianScaleMixture

• high-dimensional MixtureOfDiagNormalsSharedCovariance

Each test verifies that samples, losses, and parameter gradients remain finite.

[martinjankowiak]

thanks @jburroni! your diff changed a bunch of other files so you'll need to fix that if you want this reviewed and merged

[jburroni]

Hi @martinjankowiak!
The commits for this PR are ebd8d57 and 8384f6f.
The third one, which modified many files, was made by running make lint.

The GH actions seem to be broken, and I don't know how to proceed. What do you suggest?

[jburroni]

Hi @martinjankowiak, I've removed everything that was not strictly related to the PR.

[martinjankowiak]

@jburroni thanks! i think we'll want to wait for #3453 to merge, then rebase, and then we can review, and then we should be good to go

[martinjankowiak]

@jburroni fyi #3453 merged

[martinjankowiak]

thanks @jburroni !

[jburroni]

Thank you for your feedback @martinjankowiak

I made a mistake in my previous version. I'm sorry. I’ve now improved my changes to gaussian_scale_mixture.py.
In particular, I introduced a variable, log_gaussian_normalizer, made the code closer to the original, and used torch.erfc instead of 1-torch.erf.

I also removed manual seeding from the new tests.

[martinjankowiak]

looks good, thanks for doing this!

[martinjankowiak]

does test_mean_gradient pass with D=3? unclear to me if we're testing this branch properly. this would be an existing issue, but good to check?