diff --git a/src/pyrecest/distributions/hypersphere_subset/complex_bingham_distribution.py b/src/pyrecest/distributions/hypersphere_subset/complex_bingham_distribution.py index 05810152d..809d3c1e4 100644 --- a/src/pyrecest/distributions/hypersphere_subset/complex_bingham_distribution.py +++ b/src/pyrecest/distributions/hypersphere_subset/complex_bingham_distribution.py @@ -148,7 +148,7 @@ def pdf(self, xs): Bxs = self.B @ xs # (d, n) vals = real(einsum("ij,ij->j", conj(xs), Bxs)) # shape (n,) p = exp(self.log_norm_const + vals) - return float(p[0]) if single else p + return p[0] if single else p def sample(self, n): """Draw samples from the complex Bingham distribution. diff --git a/tests/distributions/test_complex_bingham_pytorch_autograd.py b/tests/distributions/test_complex_bingham_pytorch_autograd.py new file mode 100644 index 000000000..337c52f9c --- /dev/null +++ b/tests/distributions/test_complex_bingham_pytorch_autograd.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +import pytest + +import pyrecest.backend +from pyrecest.distributions import ComplexBinghamDistribution + +torch = pytest.importorskip("torch") + +pytestmark = pytest.mark.skipif( + pyrecest.backend.__backend_name__ != "pytorch", + reason="PyTorch backend regression", +) + + +def test_single_point_pdf_preserves_pytorch_autograd() -> None: + distribution = ComplexBinghamDistribution( + torch.tensor( + [[-3.0, 0.0], [0.0, 0.0]], + dtype=torch.complex128, + ) + ) + point = torch.tensor( + [0.5 + 0.5j, 0.5 - 0.5j], + dtype=torch.complex128, + requires_grad=True, + ) + + density = distribution.pdf(point) + + assert torch.is_tensor(density) + assert density.ndim == 0 + assert density.dtype == torch.float64 + assert density.device == point.device + assert density.requires_grad + assert torch.isfinite(density) + assert density > 0.0 + + density.backward() + + assert point.grad is not None + assert torch.all(torch.isfinite(point.grad)) + assert torch.any(point.grad != 0.0)