diff --git a/src/pyrecest/distributions/cart_prod/partially_wrapped_normal_distribution.py b/src/pyrecest/distributions/cart_prod/partially_wrapped_normal_distribution.py index cb846f9b6..0e42e7de4 100644 --- a/src/pyrecest/distributions/cart_prod/partially_wrapped_normal_distribution.py +++ b/src/pyrecest/distributions/cart_prod/partially_wrapped_normal_distribution.py @@ -31,7 +31,6 @@ tile, where, ) -from scipy.stats import multivariate_normal from ..hypertorus.hypertoroidal_wrapped_normal_distribution import ( HypertoroidalWrappedNormalDistribution, @@ -220,9 +219,10 @@ def pdf(self, xs, m: Union[int, int32, int64] = 3): axis=1, ) - # evaluate normal for all xs_wrapped - mvn = multivariate_normal(self.mu, self.C) - evals = array(mvn.pdf(xs_wrapped)) # For being compatible with all backends + # Evaluate the Gaussian factor without leaving the active backend. + evals = GaussianDistribution( + self.mu, self.C, check_validity=False + ).pdf(xs_wrapped) # sum evaluations for the wrapped dimensions summed_evals = sum(evals.reshape(-1, (2 * m + 1) ** self.bound_dim), axis=1) diff --git a/tests/distributions/test_partially_wrapped_normal_pytorch_autograd.py b/tests/distributions/test_partially_wrapped_normal_pytorch_autograd.py new file mode 100644 index 000000000..3233af2a2 --- /dev/null +++ b/tests/distributions/test_partially_wrapped_normal_pytorch_autograd.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +import pytest + +import pyrecest.backend +from pyrecest.distributions.cart_prod.partially_wrapped_normal_distribution import ( + PartiallyWrappedNormalDistribution, +) + +torch = pytest.importorskip("torch") + +pytestmark = pytest.mark.skipif( + pyrecest.backend.__backend_name__ != "pytorch", + reason="PyTorch backend regression", +) + + +def test_partially_wrapped_normal_pdf_preserves_pytorch_autograd() -> None: + dtype = torch.float64 + distribution = PartiallyWrappedNormalDistribution( + torch.tensor([0.4, -0.2], dtype=dtype), + torch.tensor([[0.8, 0.1], [0.1, 1.2]], dtype=dtype), + bound_dim=1, + ) + points = torch.tensor( + [[0.7, 0.3], [1.1, -0.5]], + dtype=dtype, + requires_grad=True, + ) + + density = distribution.pdf(points, m=1) + + assert torch.is_tensor(density) + assert density.shape == (2,) + assert density.device == points.device + assert density.dtype == points.dtype + assert density.requires_grad + assert torch.all(torch.isfinite(density)) + assert torch.all(density > 0.0) + + density.sum().backward() + + assert points.grad is not None + assert torch.all(torch.isfinite(points.grad)) + assert torch.any(points.grad != 0.0)