From 7e507b37d0088cc5becd051b21af3052e64b2a0f Mon Sep 17 00:00:00 2001 From: Florian Pfaff <6773539+FlorianPfaff@users.noreply.github.com> Date: Fri, 7 Aug 2026 01:30:27 +0800 Subject: [PATCH 1/2] Preserve backend tensors in partially wrapped normal PDFs --- .../cart_prod/partially_wrapped_normal_distribution.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) 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 cb846f9b6c..0e42e7de44 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) From b968f5d5e9c8d7af6bbc5b5c56583f794a69fedf Mon Sep 17 00:00:00 2001 From: Florian Pfaff <6773539+FlorianPfaff@users.noreply.github.com> Date: Fri, 7 Aug 2026 01:30:53 +0800 Subject: [PATCH 2/2] Test partially wrapped normal PyTorch autograd --- ...rtially_wrapped_normal_pytorch_autograd.py | 45 +++++++++++++++++++ 1 file changed, 45 insertions(+) create mode 100644 tests/distributions/test_partially_wrapped_normal_pytorch_autograd.py 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 0000000000..3233af2a23 --- /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)