Preserve PyTorch autograd in partially wrapped normal PDFs

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)

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)