From fa0b042aa7460eea88a1e4ef3e0e133c1af51560 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=A9r=C3=A9mie=20Gince?= <50332514+JeremieGince@users.noreply.github.com> Date: Fri, 19 Jun 2026 22:15:45 -0400 Subject: [PATCH 1/4] Refactor pfaffian_grad_matrix for better stability --- .../strategies/pfaffian_parlett_reid.py | 2 +- src/torch_pfaffian/strategies/strategy.py | 19 ++++- tests/test_strategies/test_strategy.py | 70 +++++++++++++++++++ 3 files changed, 87 insertions(+), 4 deletions(-) diff --git a/src/torch_pfaffian/strategies/pfaffian_parlett_reid.py b/src/torch_pfaffian/strategies/pfaffian_parlett_reid.py index 07773db..8ca854f 100644 --- a/src/torch_pfaffian/strategies/pfaffian_parlett_reid.py +++ b/src/torch_pfaffian/strategies/pfaffian_parlett_reid.py @@ -14,7 +14,7 @@ class PfaffianParlettReid(PfaffianStrategy): is the signed Pfaffian, unlike the determinant-based strategies which return only the magnitude. The backward uses the closed form ``d pf(A) / d A = (1 / 2) pf(A) (A^{-1})^T`` via the Pfaffian adjugate, with no autograd graph over the elimination. For invertible inputs this is a single - pseudo-inverse; for singular inputs (``pf == 0``) the adjugate is computed exactly from minor + inverse; for singular inputs (``pf == 0``) the adjugate is computed exactly from minor Pfaffians, so the gradient is correct everywhere (see :meth:`PfaffianStrategy.pfaffian_grad_matrix`). The input is a skew-symmetric matrix of shape ``(..., 2n, 2n)``. diff --git a/src/torch_pfaffian/strategies/strategy.py b/src/torch_pfaffian/strategies/strategy.py index 236d0c1..d2fef02 100644 --- a/src/torch_pfaffian/strategies/strategy.py +++ b/src/torch_pfaffian/strategies/strategy.py @@ -27,22 +27,35 @@ def pfaffian_grad_matrix( Gradient of the signed Pfaffian with respect to the input matrix. Uses the closed form ``d pf(A) / d A = (1 / 2) pf(A) (A^{-1})^T`` via the Pfaffian adjugate - ``pf(A) A^{-1}``. For invertible inputs the adjugate is ``pf(A) * pinv(A)`` (a single inverse); + ``pf(A) A^{-1}``. For invertible inputs the adjugate is ``pf(A) * inv(A)`` (a single inverse); for singular inputs (``pf == 0``), where that product would be ``0`` and miss the true derivative, the adjugate is recomputed exactly from minor Pfaffians via :meth:`_pfaffian_adjugate` (using ``cls``'s own forward). The minor-based path runs only on the singular batch elements, so invertible inputs keep the single cheap inverse. + The inverse uses :func:`torch.linalg.inv` (an LU factorization) rather than + :func:`torch.linalg.pinv` (an SVD). A skew-symmetric ``A`` is invertible exactly when + ``pf(A) != 0`` (since ``det(A) = pf(A)^2``), so the inverse is only ever relied upon on the + invertible elements, where the LU factorization is the correct and robust tool. The SVD-based + pseudo-inverse can fail to converge on ill-conditioned or near-repeated-singular-value inputs, + which the LU factorization does not. Because ``inv`` raises on an exactly-singular matrix, the + ``pf == 0`` elements (whose inverse is discarded anyway) are replaced by the identity before the + batched inverse so the call stays well-posed. + :param matrix: The saved input matrix of shape ``(..., n, n)``. :param pfaffian: The saved forward Pfaffian of shape ``(...,)``. :param grad_output: Gradient of the output with respect to the loss, of shape ``(...,)``. :return: Gradient of the input matrix, of shape ``(..., n, n)``. :rtype: torch.Tensor """ - adjugate = pfaffian[..., None, None] * torch.linalg.pinv(matrix) # pf(A) A^{-1}; 0 where pf == 0 singular = pfaffian == 0 + dimension = matrix.shape[-1] + # inv (LU) raises on exactly-singular matrices, so the pf == 0 elements are replaced by the + # identity before the batched inverse; their adjugate is overwritten from minor Pfaffians below. + identity = torch.eye(dimension, dtype=matrix.dtype, device=matrix.device).expand_as(matrix) + safe_matrix = torch.where(singular[..., None, None], identity, matrix) # (..., n, n) + adjugate = pfaffian[..., None, None] * torch.linalg.inv(safe_matrix) # pf(A) A^{-1}; 0 where pf == 0 if bool(singular.any()): - dimension = matrix.shape[-1] flat_matrix = matrix.reshape(-1, dimension, dimension) flat_adjugate = adjugate.reshape(-1, dimension, dimension) singular_index = singular.reshape(-1).nonzero(as_tuple=True)[0] diff --git a/tests/test_strategies/test_strategy.py b/tests/test_strategies/test_strategy.py index f8644e5..108bae3 100644 --- a/tests/test_strategies/test_strategy.py +++ b/tests/test_strategies/test_strategy.py @@ -58,3 +58,73 @@ def test_grad_matrix_uses_adjugate_on_singular(self): # d pf / d A_{01} = a_{23} = 1, so the gradient of the singular element is nonzero. assert torch.isfinite(result).all() assert result[0].abs().sum() > 0 + + def test_grad_matrix_mixed_batch_matches_closed_forms(self): + # A batch mixing a singular (pf=0) and an invertible element must not raise (a plain + # torch.linalg.inv would, on the singular pivot) and each element must match its exact + # closed form: the invertible one via the inverse, the singular one via the minor adjugate. + singular = torch.zeros(4, 4, dtype=torch.float64) + singular[2, 3] = 1.0 + singular[3, 2] = -1.0 + invertible = _random_skew(4, seed=2) + matrix = torch.stack([singular, invertible]) + pfaffian = PfaffianParlettReid.forward(matrix) + grad_output = torch.tensor([1.3, -0.7], dtype=torch.float64) + result = PfaffianParlettReid.pfaffian_grad_matrix(matrix, pfaffian, grad_output) + + expected_invertible = torch.einsum( + "ij->ji", 0.5 * grad_output[1] * pfaffian[1] * torch.linalg.inv(invertible) + ) + minor_adjugate = PfaffianParlettReid._pfaffian_adjugate(singular[None])[0] + expected_singular = torch.einsum("ij->ji", 0.5 * grad_output[0] * minor_adjugate) + torch.testing.assert_close( + result[1], expected_invertible, atol=ATOL_MATRIX_COMPARISON, rtol=RTOL_MATRIX_COMPARISON + ) + torch.testing.assert_close( + result[0], expected_singular, atol=ATOL_MATRIX_COMPARISON, rtol=RTOL_MATRIX_COMPARISON + ) + + def test_grad_matrix_all_singular_batch_is_finite_and_exact(self): + # Every element singular: the inverse path is entirely bypassed by the minor adjugate. + singular = torch.zeros(4, 4, dtype=torch.float64) + singular[2, 3] = 1.0 + singular[3, 2] = -1.0 + matrix = torch.stack([singular, singular.clone()]) + pfaffian = PfaffianParlettReid.forward(matrix) + assert bool((pfaffian == 0).all()) + grad_output = torch.tensor([1.0, 2.0], dtype=torch.float64) + result = PfaffianParlettReid.pfaffian_grad_matrix(matrix, pfaffian, grad_output) + assert torch.isfinite(result).all() + minor_adjugate = PfaffianParlettReid._pfaffian_adjugate(matrix) + expected = torch.einsum("...,...ij->...ji", 0.5 * grad_output, minor_adjugate) + torch.testing.assert_close(result, expected, atol=ATOL_MATRIX_COMPARISON, rtol=RTOL_MATRIX_COMPARISON) + + def test_grad_matrix_does_not_use_svd(self, monkeypatch): + # Regression for the SVD-non-convergence crash: the gradient must not route through + # torch.linalg.pinv (SVD) anymore. Force pinv to fail and confirm the gradient still computes. + def _fail(*args, **kwargs): + raise AssertionError("torch.linalg.pinv (SVD) must not be used by pfaffian_grad_matrix") + + monkeypatch.setattr(torch.linalg, "pinv", _fail) + singular = torch.zeros(4, 4, dtype=torch.float64) + singular[2, 3] = 1.0 + singular[3, 2] = -1.0 + matrix = torch.stack([singular, _random_skew(4, seed=3)]) + pfaffian = PfaffianParlettReid.forward(matrix) + grad_output = torch.ones(2, dtype=torch.float64) + result = PfaffianParlettReid.pfaffian_grad_matrix(matrix, pfaffian, grad_output) + assert torch.isfinite(result).all() + + def test_grad_matrix_ill_conditioned_invertible_is_finite_and_exact(self): + # The bug surfaced on ill-conditioned skew inputs where the SVD pseudo-inverse fails to + # converge. The LU inverse stays finite; the gradient must match the inverse closed form. + scales = torch.tensor([1e8, 1e-8, 1e6, 1e-6], dtype=torch.float64) + blocks = [torch.tensor([[0.0, scale], [-scale, 0.0]], dtype=torch.float64) for scale in scales] + matrix = torch.block_diag(*blocks) # invertible, condition number ~1e16 + pfaffian = PfaffianParlettReid.forward(matrix) + assert pfaffian != 0.0 + grad_output = torch.tensor(1.0, dtype=torch.float64) + result = PfaffianParlettReid.pfaffian_grad_matrix(matrix, pfaffian, grad_output) + expected = torch.einsum("ij->ji", 0.5 * grad_output * pfaffian * torch.linalg.inv(matrix)) + assert torch.isfinite(result).all() + torch.testing.assert_close(result, expected, atol=ATOL_MATRIX_COMPARISON, rtol=RTOL_MATRIX_COMPARISON) From 07b9655502b31dfac51f56941583157a713c7a35 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=A9r=C3=A9mie=20Gince?= <50332514+JeremieGince@users.noreply.github.com> Date: Fri, 19 Jun 2026 23:04:27 -0400 Subject: [PATCH 2/4] Enhance backward methods for Pfaffian calculations to handle singular blocks and improve stability --- .../strategies/pfaffian_block_det.py | 68 +++++++++++++++---- .../strategies/pfaffian_fdbpf.py | 44 ++++++++---- .../strategies/pfaffian_parlett_reid.py | 14 ++-- src/torch_pfaffian/strategies/strategy.py | 15 ++-- .../test_pfaffian_block_det.py | 50 ++++++++++++++ tests/test_strategies/test_pfaffian_fdbpf.py | 49 +++++++++++++ 6 files changed, 202 insertions(+), 38 deletions(-) diff --git a/src/torch_pfaffian/strategies/pfaffian_block_det.py b/src/torch_pfaffian/strategies/pfaffian_block_det.py index dd083ee..d29b646 100644 --- a/src/torch_pfaffian/strategies/pfaffian_block_det.py +++ b/src/torch_pfaffian/strategies/pfaffian_block_det.py @@ -29,14 +29,20 @@ def forward(matrix: torch.Tensor): return pf @staticmethod - def backward(ctx: torch.autograd.function.BackwardCFunction, grad_output): + def backward(ctx: torch.autograd.function.BackwardCFunction, grad_output: torch.Tensor) -> torch.Tensor | None: r""" Compute the gradient of the Pfaffian with respect to the full input matrix. The forward depends only on the upper-right block ``block = matrix[..., :n, n:]`` through - ``pf = c * det(block)``. By Jacobi's formula ``d det(block) / d block = det(block) * (block^{-1})^T``, - so ``d pf / d block = pf * (block^{-1})^T``. The gradient is therefore zero everywhere except in - the upper-right block. + ``pf = c * det(block)``. By Jacobi's formula ``d det(block) / d block = det(block) (block^{-1})^T``, + so ``d pf / d block = c * adj(block)^T`` where ``adj(block) = det(block) block^{-1}`` is the + classical adjugate. The gradient is therefore zero everywhere except in the upper-right block. + + For an invertible block this is the single inverse ``pf * (block^{-1})^T``. For a singular block + (``pf == 0``) that product would be ``0 * inf`` and :func:`torch.linalg.inv` raises, yet the true + derivative ``c * adj(block)^T`` is finite and generally nonzero. The adjugate is then computed + exactly from the cofactors via :meth:`_cofactor_matrix`, on the singular batch elements only so + invertible inputs keep the single cheap inverse. :param ctx: Context holding the saved input matrix and the forward Pfaffian. :param grad_output: Gradient of the output with respect to the loss. @@ -44,16 +50,54 @@ def backward(ctx: torch.autograd.function.BackwardCFunction, grad_output): :rtype: torch.Tensor | None """ matrix, pf = cast("tuple[torch.Tensor, torch.Tensor]", ctx.saved_tensors) - grad_matrix = None - if ctx.needs_input_grad[0]: - n = matrix.shape[-1] // 2 - sub_matrix = matrix[..., :n, n:] # (..., n, n) upper-right block - inverse_transpose = torch.linalg.inv(sub_matrix).transpose(-1, -2) - grad_block = grad_output[..., None, None] * pf[..., None, None] * inverse_transpose - grad_matrix = torch.zeros_like(matrix) - grad_matrix[..., :n, n:] = grad_block + if not ctx.needs_input_grad[0]: + return None + n = matrix.shape[-1] // 2 + block = matrix[..., :n, n:] # (..., n, n) upper-right block + constant = (-1) ** (n * (n - 1) // 2) + singular = pf == 0 + if bool(singular.any()): + identity = torch.eye(n, dtype=matrix.dtype, device=matrix.device).expand_as(block) + safe_block = torch.where(singular[..., None, None], identity, block) + adjugate_transpose = pf[..., None, None] * torch.linalg.inv(safe_block).transpose(-1, -2) + flat_block = block.reshape(-1, n, n) + flat_adjugate = adjugate_transpose.reshape(-1, n, n) + singular_index = singular.reshape(-1).nonzero(as_tuple=True)[0] + cofactor = constant * PfaffianBlockDet._cofactor_matrix(flat_block.index_select(0, singular_index)) + flat_adjugate = flat_adjugate.index_copy(0, singular_index, cofactor.to(flat_adjugate.dtype)) + adjugate_transpose = flat_adjugate.reshape_as(block) + else: + adjugate_transpose = pf[..., None, None] * torch.linalg.inv(block).transpose(-1, -2) + grad_matrix = torch.zeros_like(matrix) + grad_matrix[..., :n, n:] = grad_output[..., None, None] * adjugate_transpose return grad_matrix + @staticmethod + def _cofactor_matrix(blocks: torch.Tensor) -> torch.Tensor: + r""" + Cofactor matrix ``C`` of a batch of square blocks, with ``C_{ij} = (-1)^{i+j} det(B^{(ij)})`` + where ``B^{(ij)}`` is ``B`` with row ``i`` and column ``j`` removed. + + This equals the transpose of the classical adjugate ``adj(B) = det(B) B^{-1}``, so it stays + finite when ``B`` is singular and an inverse would fail. The minor determinants are batched, so + the cost is ``n^2`` determinant evaluations regardless of the batch size. + + :param blocks: Square blocks of shape ``(m, n, n)``. + :return: The cofactor matrix of shape ``(m, n, n)``. + :rtype: torch.Tensor + """ + n = blocks.shape[-1] + cofactor = torch.zeros_like(blocks) + indices = torch.arange(n, device=blocks.device) + for row in range(n): + kept_rows = indices[indices != row] + for column in range(n): + kept_columns = indices[indices != column] + minor = blocks.index_select(-2, kept_rows).index_select(-1, kept_columns) # (m, n-1, n-1) + sign = 1.0 if (row + column) % 2 == 0 else -1.0 + cofactor[..., row, column] = sign * torch.linalg.det(minor) + return cofactor + @classmethod def _pfaffian_adjugate(cls, matrices: torch.Tensor) -> torch.Tensor: # This strategy's forward only accepts block-antidiagonal matrices, but the Pfaffian minors diff --git a/src/torch_pfaffian/strategies/pfaffian_fdbpf.py b/src/torch_pfaffian/strategies/pfaffian_fdbpf.py index 1695251..732d57f 100644 --- a/src/torch_pfaffian/strategies/pfaffian_fdbpf.py +++ b/src/torch_pfaffian/strategies/pfaffian_fdbpf.py @@ -23,18 +23,38 @@ def forward(matrix: torch.Tensor): return pf @staticmethod - def backward(ctx: torch.autograd.function.BackwardCFunction, grad_output): + def backward(ctx: torch.autograd.function.BackwardCFunction, grad_output: torch.Tensor) -> torch.Tensor | None: r""" - - ..math: - \frac{\partial \text{pf}(A)}{\partial A_{ij}} = \frac{\text{pf}(A)}{2} A^{-1}_{ji} - - :param ctx: Context - :param grad_output: Gradient of the output - :return: Gradient of the input + Gradient of the Pfaffian magnitude with respect to the input matrix. + + .. math:: + \frac{\partial |\text{pf}(A)|}{\partial A_{ij}} = \frac{|\text{pf}(A)|}{2} (A^{-1})_{ji} + + The inverse uses :func:`torch.linalg.inv` (an LU factorization) rather than + :func:`torch.linalg.pinv` (an SVD), which can fail to converge on ill-conditioned or + near-repeated-singular-value inputs. The forward floors the magnitude at ``sqrt(EPSILON)``, so + singular elements (``pf`` at that floor, and the always-singular odd-dimensional inputs whose + ``pf`` is ``0``) are replaced by the identity before the batched inverse: ``inv`` raises on an + exactly-singular matrix, and there the gradient is negligible anyway since it scales with ``pf``. + + :param ctx: Context holding the saved input matrix and the forward Pfaffian magnitude. + :param grad_output: Gradient of the output with respect to the loss. + :return: Gradient of the input matrix, or ``None`` when the input does not require grad. + :rtype: torch.Tensor | None """ matrix, pf = cast("tuple[torch.Tensor, torch.Tensor]", ctx.saved_tensors) - grad_matrix = None - if ctx.needs_input_grad[0]: - grad_matrix = torch.einsum("...,...ij->...ji", 0.5 * grad_output * pf, torch.linalg.pinv(matrix)) - return grad_matrix + if not ctx.needs_input_grad[0]: + return None + # The forward clamps the radicand at EPSILON, so a singular element has pf exactly at the floor + # sqrt(EPSILON) (or 0 for odd dimensions). Computing the floor with the same dtype and sqrt as + # the forward makes the comparison exact rather than dependent on a hand-written threshold. + singular_floor = matrix.new_tensor(PfaffianFDBPf.EPSILON).sqrt() + singular = pf <= singular_floor + if bool(singular.any()): + dimension = matrix.shape[-1] + identity = torch.eye(dimension, dtype=matrix.dtype, device=matrix.device).expand_as(matrix) + safe_matrix = torch.where(singular[..., None, None], identity, matrix) + inverse = torch.linalg.inv(safe_matrix) + else: + inverse = torch.linalg.inv(matrix) + return torch.einsum("...,...ij->...ji", 0.5 * grad_output * pf, inverse) diff --git a/src/torch_pfaffian/strategies/pfaffian_parlett_reid.py b/src/torch_pfaffian/strategies/pfaffian_parlett_reid.py index 8ca854f..4a51256 100644 --- a/src/torch_pfaffian/strategies/pfaffian_parlett_reid.py +++ b/src/torch_pfaffian/strategies/pfaffian_parlett_reid.py @@ -39,22 +39,20 @@ def forward(matrix: torch.Tensor) -> torch.Tensor: for column in range(0, dimension - 2, 2): sub_column = working[:, column + 2 :, column].abs() # (batch, 2n - column - 2) - pivot_relative = sub_column.argmax(dim=1) # (batch,) + pivot_magnitude, pivot_relative = sub_column.max(dim=1) # (batch,), (batch,) pivot_row = pivot_relative + column + 2 # (batch,) - pivot_magnitude = sub_column.gather(1, pivot_relative[:, None]).squeeze(1) # (batch,) pivot_condition = pivot_magnitude > working[:, column + 1, column].abs() # (batch,) mask = pivot_condition[:, None] # (batch, 1) - # Congruence swap of rows/columns column+1 <-> pivot_row where pivoting helps. row_fixed = working[:, column + 1, :].clone() - row_pivot = working[batch_index, pivot_row, :].clone() + row_pivot = working[batch_index, pivot_row, :] working[:, column + 1, :] = torch.where(mask, row_pivot, row_fixed) working[batch_index, pivot_row, :] = torch.where(mask, row_fixed, row_pivot) col_fixed = working[:, :, column + 1].clone() - col_pivot = working[batch_index, :, pivot_row].clone() + col_pivot = working[batch_index, :, pivot_row] working[:, :, column + 1] = torch.where(mask, col_pivot, col_fixed) working[batch_index, :, pivot_row] = torch.where(mask, col_fixed, col_pivot) - sign = sign * torch.where(pivot_condition, -torch.ones_like(sign), torch.ones_like(sign)) + sign = torch.where(pivot_condition, -sign, sign) pivot_value = working[:, column + 1, column] # (batch,) zero_pivot = pivot_value.abs() < epsilon @@ -62,8 +60,8 @@ def forward(matrix: torch.Tensor) -> torch.Tensor: safe_pivot = torch.where(zero_pivot, torch.ones_like(pivot_value), pivot_value) tau = working[:, column + 2 :, column] / safe_pivot[:, None] # (batch, 2n - column - 2) column_next = working[:, column + 2 :, column + 1] # (batch, 2n - column - 2) - update = torch.einsum("bi,bj->bij", tau, column_next) - torch.einsum("bi,bj->bij", column_next, tau) - working[:, column + 2 :, column + 2 :] = working[:, column + 2 :, column + 2 :] + update + outer = tau[:, :, None] * column_next[:, None, :] # (batch, m, m) + working[:, column + 2 :, column + 2 :].add_(outer - outer.transpose(-1, -2)) working[:, column + 2 :, column] = 0 working[:, column, column + 2 :] = 0 working[:, column + 2 :, column + 1] = 0 diff --git a/src/torch_pfaffian/strategies/strategy.py b/src/torch_pfaffian/strategies/strategy.py index d2fef02..171cb8e 100644 --- a/src/torch_pfaffian/strategies/strategy.py +++ b/src/torch_pfaffian/strategies/strategy.py @@ -50,12 +50,15 @@ def pfaffian_grad_matrix( """ singular = pfaffian == 0 dimension = matrix.shape[-1] - # inv (LU) raises on exactly-singular matrices, so the pf == 0 elements are replaced by the - # identity before the batched inverse; their adjugate is overwritten from minor Pfaffians below. - identity = torch.eye(dimension, dtype=matrix.dtype, device=matrix.device).expand_as(matrix) - safe_matrix = torch.where(singular[..., None, None], identity, matrix) # (..., n, n) - adjugate = pfaffian[..., None, None] * torch.linalg.inv(safe_matrix) # pf(A) A^{-1}; 0 where pf == 0 - if bool(singular.any()): + any_singular = bool(singular.any()) + if any_singular: + identity = torch.eye(dimension, dtype=matrix.dtype, device=matrix.device).expand_as(matrix) + safe_matrix = torch.where(singular[..., None, None], identity, matrix) # (..., n, n) + inverse = torch.linalg.inv(safe_matrix) + else: + inverse = torch.linalg.inv(matrix) + adjugate = pfaffian[..., None, None] * inverse # pf(A) A^{-1}; 0 where pf == 0 + if any_singular: flat_matrix = matrix.reshape(-1, dimension, dimension) flat_adjugate = adjugate.reshape(-1, dimension, dimension) singular_index = singular.reshape(-1).nonzero(as_tuple=True)[0] diff --git a/tests/test_strategies/test_pfaffian_block_det.py b/tests/test_strategies/test_pfaffian_block_det.py index a67800f..c7d9881 100644 --- a/tests/test_strategies/test_pfaffian_block_det.py +++ b/tests/test_strategies/test_pfaffian_block_det.py @@ -81,6 +81,56 @@ class _Context: assert PfaffianBlockDet.backward(_Context(), torch.ones_like(pfaffian)) is None + @pytest.mark.parametrize("size", _BLOCK_SIZES) + def test_cofactor_matrix_matches_det_inverse_on_invertible(self, size): + # On invertible blocks the cofactor matrix must equal the classical adjugate transpose + # det(B) * (B^{-1})^T, the relation the singular gradient relies on. + blocks = torch.tensor(_RNG.random((4, size, size))) + torch.eye(size) + cofactor = PfaffianBlockDet._cofactor_matrix(blocks) + expected = torch.linalg.det(blocks)[..., None, None] * torch.linalg.inv(blocks).transpose(-1, -2) + torch.testing.assert_close(cofactor, expected, atol=ATOL_MATRIX_COMPARISON, rtol=RTOL_MATRIX_COMPARISON) + + def test_backward_singular_block_is_finite_and_exact(self): + # A singular upper-right block (pf=0) used to crash the inverse-based backward. The gradient is + # c * adj(B)^T (finite, generally nonzero); verify it against an independent numpy cofactor. + block = np.array([[1.0, 2.0, 3.0], [2.0, 4.0, 6.0], [1.0, 0.0, 5.0]]) # rows 0,1 dependent -> det 0 + matrix = _block_antidiagonal(block).requires_grad_(True) + pfaffian = PfaffianBlockDet.apply(matrix) + assert pfaffian.item() == 0.0 + pfaffian.backward() + assert torch.isfinite(matrix.grad).all() + + size = block.shape[0] + constant = (-1) ** (size * (size - 1) // 2) + cofactor = np.array( + [ + [(-1) ** (i + j) * np.linalg.det(np.delete(np.delete(block, i, 0), j, 1)) for j in range(size)] + for i in range(size) + ] + ) + expected_block = torch.tensor(constant * cofactor) + torch.testing.assert_close( + matrix.grad[:size, size:], expected_block, atol=ATOL_MATRIX_COMPARISON, rtol=RTOL_MATRIX_COMPARISON + ) + assert matrix.grad[:size, :size].abs().max() == 0 # gradient is confined to the upper-right block + assert matrix.grad[size:, :].abs().max() == 0 + + def test_backward_mixed_singular_invertible_batch_matches_inverse(self): + # A batch mixing a singular block with an invertible one must not raise, and the invertible + # element's gradient must match the inverse-based closed form. + singular_block = torch.zeros(3, 3, dtype=torch.float64) + invertible_block = torch.tensor(_RNG.random((3, 3))) + torch.eye(3) + matrix = torch.stack( + [_block_antidiagonal(singular_block), _block_antidiagonal(invertible_block)] + ).requires_grad_(True) + PfaffianBlockDet.apply(matrix).sum().backward() + assert torch.isfinite(matrix.grad).all() + constant = (-1) ** (3 * (3 - 1) // 2) + expected = constant * torch.linalg.det(invertible_block) * torch.linalg.inv(invertible_block).transpose(-1, -2) + torch.testing.assert_close( + matrix.grad[1, :3, 3:], expected, atol=ATOL_MATRIX_COMPARISON, rtol=RTOL_MATRIX_COMPARISON + ) + def test_adjugate_override_delegates_to_parlett_reid(self): # PfaffianBlockDet.forward is block-only, so its adjugate must defer to the general strategy. from torch_pfaffian.strategies.pfaffian_parlett_reid import PfaffianParlettReid diff --git a/tests/test_strategies/test_pfaffian_fdbpf.py b/tests/test_strategies/test_pfaffian_fdbpf.py index 3f7246f..b6e4f75 100644 --- a/tests/test_strategies/test_pfaffian_fdbpf.py +++ b/tests/test_strategies/test_pfaffian_fdbpf.py @@ -101,3 +101,52 @@ class _Context: needs_input_grad = (False,) assert PfaffianFDBPf.backward(_Context(), torch.ones_like(pfaffian)) is None + + def test_backward_does_not_use_svd(self, monkeypatch): + # Regression: the magnitude gradient must not route through torch.linalg.pinv (SVD), whose + # iterative solver can fail to converge on ill-conditioned inputs. Force pinv to fail. + def _fail(*args, **kwargs): + raise AssertionError("torch.linalg.pinv (SVD) must not be used by PfaffianFDBPf.backward") + + monkeypatch.setattr(torch.linalg, "pinv", _fail) + matrix = torch.tensor(_skew(_RNG.random((6, 6))), requires_grad=True) + PfaffianFDBPf.apply(matrix).backward() + assert torch.isfinite(matrix.grad).all() + + def test_backward_odd_dimension_gradient_is_finite_zero(self): + # Odd-dimensional skew matrices are always singular (inv would raise / return garbage); the + # magnitude is 0 there, so the gradient must be exactly zero and finite, never NaN. + matrix = torch.tensor(_skew(_RNG.random((5, 5))), requires_grad=True) + PfaffianFDBPf.apply(matrix).backward() + assert torch.isfinite(matrix.grad).all() + torch.testing.assert_close(matrix.grad, torch.zeros_like(matrix)) + + def test_backward_mixed_singular_invertible_batch_matches_inverse(self): + # A batch mixing an exactly-singular element (pf at the floor) with an invertible one must not + # raise, and the invertible element's gradient must match the inverse-based closed form. + singular = torch.zeros(4, 4, dtype=torch.float64) + singular[2, 3] = 1.0 + singular[3, 2] = -1.0 + invertible = torch.tensor(_skew(_RNG.random((4, 4)))) + matrix = torch.stack([singular, invertible]).requires_grad_(True) + magnitude = PfaffianFDBPf.apply(matrix) + magnitude.sum().backward() + assert torch.isfinite(matrix.grad).all() + expected_invertible = torch.einsum( + "ij->ji", 0.5 * magnitude[1].detach() * torch.linalg.inv(invertible) + ) + torch.testing.assert_close( + matrix.grad[1], expected_invertible, atol=ATOL_MATRIX_COMPARISON, rtol=RTOL_MATRIX_COMPARISON + ) + + def test_backward_ill_conditioned_is_finite_and_matches_inverse(self): + # Ill-conditioned invertible input (the regime where the SVD pseudo-inverse fails to converge): + # the LU inverse stays finite and the gradient matches the inverse closed form. + scales = torch.tensor([1e8, 1e-8, 1e6, 1e-6], dtype=torch.float64) + blocks = [torch.tensor([[0.0, scale], [-scale, 0.0]], dtype=torch.float64) for scale in scales] + matrix = torch.block_diag(*blocks).requires_grad_(True) + magnitude = PfaffianFDBPf.apply(matrix) + magnitude.backward() + expected = torch.einsum("ij->ji", 0.5 * magnitude.detach() * torch.linalg.inv(matrix.detach())) + assert torch.isfinite(matrix.grad).all() + torch.testing.assert_close(matrix.grad, expected, atol=ATOL_MATRIX_COMPARISON, rtol=RTOL_MATRIX_COMPARISON) From 7606fc2b44b702d66375203c836a0b48ca33ae8e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=A9r=C3=A9mie=20Gince?= <50332514+JeremieGince@users.noreply.github.com> Date: Fri, 19 Jun 2026 23:10:11 -0400 Subject: [PATCH 3/4] Refactor tensor operations in tests for improved readability --- tests/test_strategies/test_pfaffian_fdbpf.py | 4 +--- tests/test_strategies/test_strategy.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/test_strategies/test_pfaffian_fdbpf.py b/tests/test_strategies/test_pfaffian_fdbpf.py index b6e4f75..8a4f25c 100644 --- a/tests/test_strategies/test_pfaffian_fdbpf.py +++ b/tests/test_strategies/test_pfaffian_fdbpf.py @@ -132,9 +132,7 @@ def test_backward_mixed_singular_invertible_batch_matches_inverse(self): magnitude = PfaffianFDBPf.apply(matrix) magnitude.sum().backward() assert torch.isfinite(matrix.grad).all() - expected_invertible = torch.einsum( - "ij->ji", 0.5 * magnitude[1].detach() * torch.linalg.inv(invertible) - ) + expected_invertible = torch.einsum("ij->ji", 0.5 * magnitude[1].detach() * torch.linalg.inv(invertible)) torch.testing.assert_close( matrix.grad[1], expected_invertible, atol=ATOL_MATRIX_COMPARISON, rtol=RTOL_MATRIX_COMPARISON ) diff --git a/tests/test_strategies/test_strategy.py b/tests/test_strategies/test_strategy.py index 108bae3..0519954 100644 --- a/tests/test_strategies/test_strategy.py +++ b/tests/test_strategies/test_strategy.py @@ -72,9 +72,7 @@ def test_grad_matrix_mixed_batch_matches_closed_forms(self): grad_output = torch.tensor([1.3, -0.7], dtype=torch.float64) result = PfaffianParlettReid.pfaffian_grad_matrix(matrix, pfaffian, grad_output) - expected_invertible = torch.einsum( - "ij->ji", 0.5 * grad_output[1] * pfaffian[1] * torch.linalg.inv(invertible) - ) + expected_invertible = torch.einsum("ij->ji", 0.5 * grad_output[1] * pfaffian[1] * torch.linalg.inv(invertible)) minor_adjugate = PfaffianParlettReid._pfaffian_adjugate(singular[None])[0] expected_singular = torch.einsum("ij->ji", 0.5 * grad_output[0] * minor_adjugate) torch.testing.assert_close( From f23f09571fdc32f46f15f58c9170cde993ea9d2d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=A9r=C3=A9mie=20Gince?= <50332514+JeremieGince@users.noreply.github.com> Date: Mon, 22 Jun 2026 10:10:50 -0400 Subject: [PATCH 4/4] Add support for complex inputs in Pfaffian calculations and enhance versioning in CI --- .github/workflows/build_dist.yml | 103 ++++++++-- .gitignore | 1 + rust/Cargo.lock | 48 +++++ rust/Cargo.toml | 4 +- rust/src/lib.rs | 190 ++++++++++++++++-- src/torch_pfaffian/__init__.py | 7 +- .../strategies/pfaffian_rust_parlett_reid.py | 41 +++- src/torch_pfaffian/strategies/strategy.py | 7 +- .../test_pfaffian_rust_parlett_reid.py | 86 ++++++++ tests/test_torch_pfaffian.py | 131 ++++++++++++ 10 files changed, 574 insertions(+), 44 deletions(-) diff --git a/.github/workflows/build_dist.yml b/.github/workflows/build_dist.yml index 3e52399..10f2b25 100644 --- a/.github/workflows/build_dist.yml +++ b/.github/workflows/build_dist.yml @@ -15,10 +15,42 @@ permissions: id-token: write jobs: - Build-Dist: - name: Build dist + Version: + name: Compute next version runs-on: ubuntu-latest + outputs: + new_version: ${{ steps.version.outputs.new_tag }} + steps: + - uses: actions/checkout@v3 + + - name: Gather new package version + id: version + uses: anothrNick/github-tag-action@1.61.0 + env: + GITHUB_TOKEN: ${{secrets.GITHUB_TOKEN}} + WITH_V: false + DEFAULT_BUMP: patch + DRY_RUN: true + Build-Wheels: + name: Build wheel (${{ matrix.target }}) + needs: Version + strategy: + fail-fast: false + matrix: + include: + - os: ubuntu-latest + target: linux + - os: windows-latest + target: windows + - os: macos-13 + target: macos-x86_64 + - os: macos-14 + target: macos-arm64 + runs-on: ${{ matrix.os }} + defaults: + run: + shell: bash steps: - uses: actions/checkout@v3 - name: Set up Python 3.11 @@ -37,30 +69,69 @@ jobs: uv venv .venv uv sync --locked --dev --extra cpu - - name: Gather new package version - id: version - uses: anothrNick/github-tag-action@1.61.0 - env: - GITHUB_TOKEN: ${{secrets.GITHUB_TOKEN}} - WITH_V: false - DEFAULT_BUMP: patch - DRY_RUN: true - - - name: Bump package version in pyproject.toml - run: | - uv version ${{steps.version.outputs.new_tag}} - uv lock + - name: Stamp build version + run: uv version ${{ needs.Version.outputs.new_version }} - name: Install Rust toolchain uses: dtolnay/rust-toolchain@stable - - name: Build dist + - name: Build manylinux wheel and sdist + if: matrix.target == 'linux' run: | rm -f dist/*.whl dist/*.tar.gz uv run maturin build --release --out dist --compatibility manylinux2014 --zig uv run maturin sdist --out dist uv run twine check dist/* + - name: Build native wheel + if: matrix.target != 'linux' + run: | + uv run maturin build --release --out dist + uv run twine check dist/* + + - name: Upload distribution artifacts + uses: actions/upload-artifact@v4 + with: + name: dist-${{ matrix.target }} + path: dist/* + + Publish: + name: Release and publish + needs: [Version, Build-Wheels] + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - name: Set up Python 3.11 + uses: actions/setup-python@v3 + with: + python-version: "3.11" + + - name: Install uv + uses: astral-sh/setup-uv@v6 + with: + version: "0.9.2" + python-version: "3.11" + enable-cache: true + - name: Install dependencies + run: | + uv venv .venv + uv sync --locked --dev --extra cpu + + - name: Download all distribution artifacts + uses: actions/download-artifact@v4 + with: + path: dist + pattern: dist-* + merge-multiple: true + + - name: Check distributions + run: uv run twine check dist/* + + - name: Bump package version in pyproject.toml + run: | + uv version ${{ needs.Version.outputs.new_version }} + uv lock + - name: Commit updated pyproject.toml run: | git config --local user.email "action@github.com" diff --git a/.gitignore b/.gitignore index 59da262..cc01194 100644 --- a/.gitignore +++ b/.gitignore @@ -113,6 +113,7 @@ __pycache__/ # C extensions *.so +*.pdb # Distribution / packaging .Python diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 37519ad..a1192d4 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -39,12 +39,30 @@ version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + [[package]] name = "either" version = "1.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "num-traits", + "zerocopy", +] + [[package]] name = "heck" version = "0.5.0" @@ -66,6 +84,12 @@ version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +[[package]] +name = "libm" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" + [[package]] name = "matrixmultiply" version = "0.3.10" @@ -125,6 +149,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" dependencies = [ "autocfg", + "libm", ] [[package]] @@ -133,6 +158,7 @@ version = "0.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b94caae805f998a07d33af06e6a3891e38556051b8045c615470a71590e13e78" dependencies = [ + "half", "libc", "ndarray", "num-complex", @@ -303,7 +329,9 @@ checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" name = "torch_pfaffian_rust" version = "0.0.1" dependencies = [ + "half", "ndarray", + "num-complex", "num-traits", "numpy", "pyo3", @@ -321,3 +349,23 @@ name = "unindent" version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" + +[[package]] +name = "zerocopy" +version = "0.8.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 6b8deef..8302fa8 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -9,7 +9,9 @@ crate-type = ["cdylib"] [dependencies] pyo3 = { version = "0.23", features = ["abi3-py310", "extension-module"] } -numpy = "0.23" +numpy = { version = "0.23", features = ["half"] } ndarray = "0.16" num-traits = "0.2" +num-complex = "0.4" +half = { version = "2", features = ["num-traits"] } rayon = "1" diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 0fd6912..29b3193 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -1,8 +1,11 @@ -use num_traits::Float; +use half::f16; +use num_complex::Complex; +use num_traits::{One, Zero}; use numpy::ndarray::{Array1, Array2, Axis}; -use numpy::{IntoPyArray, PyArray1, PyReadonlyArray3}; +use numpy::{Complex32, Complex64, IntoPyArray, PyArray1, PyReadonlyArray3}; use pyo3::prelude::*; use rayon::prelude::*; +use std::ops::{Add, Div, Mul, Neg, Sub}; const PIVOT_EPSILON: f64 = 1e-30; @@ -10,13 +13,65 @@ const PIVOT_EPSILON: f64 = 1e-30; /// batches run serially to avoid thread-pool overhead. const PARALLEL_BATCH_THRESHOLD: usize = 8; +/// Magnitude used to rank pivots and to test for a numerically zero pivot. +/// +/// For real scalars this is the absolute value; for complex scalars it is the modulus. The pivot +/// choice and the singularity test only ever compare magnitudes, so the elimination itself stays +/// generic over real and complex precisions while the algebra runs in the native (possibly complex) +/// type. +trait Magnitude { + fn magnitude(&self) -> f64; +} + +impl Magnitude for f16 { + fn magnitude(&self) -> f64 { + self.to_f64().abs() + } +} + +impl Magnitude for f32 { + fn magnitude(&self) -> f64 { + self.abs() as f64 + } +} + +impl Magnitude for f64 { + fn magnitude(&self) -> f64 { + self.abs() + } +} + +impl Magnitude for Complex { + fn magnitude(&self) -> f64 { + self.norm() as f64 + } +} + +impl Magnitude for Complex { + fn magnitude(&self) -> f64 { + self.norm() + } +} + /// Signed Pfaffian of a single skew-symmetric matrix via Parlett-Reid elimination. /// -/// Generic over the floating precision so the same algorithm serves ``f32`` and ``f64`` inputs. -/// The matrix is copied into a flat row-major buffer so the hot rank-2 Schur update runs over a +/// Generic over the scalar type so the same algorithm serves ``f32``/``f64`` and their complex +/// counterparts: the arithmetic runs in the native type while pivoting compares magnitudes. The +/// matrix is copied into a flat row-major buffer so the hot rank-2 Schur update runs over a /// contiguous row slice, which the compiler can auto-vectorize (far cheaper than per-element /// strided ndarray indexing). -fn pfaffian_one(matrix: Array2) -> T { +fn pfaffian_one(matrix: Array2) -> T +where + T: Copy + + Zero + + One + + Magnitude + + Neg + + Add + + Sub + + Mul + + Div, +{ let dimension = matrix.nrows(); if dimension % 2 == 1 { return T::zero(); @@ -25,21 +80,20 @@ fn pfaffian_one(matrix: Array2) -> T { return T::one(); } let mut data: Vec = matrix.iter().copied().collect(); // row-major, len dimension * dimension - let epsilon = T::from(PIVOT_EPSILON).unwrap(); let mut sign = T::one(); let mut column = 0usize; while column + 2 < dimension { - // Partial pivoting: largest |data[row, column]| for row > column + 1. + // Partial pivoting: largest magnitude data[row, column] for row > column + 1. let mut pivot_row = column + 2; - let mut best = data[(column + 2) * dimension + column].abs(); + let mut best = data[(column + 2) * dimension + column].magnitude(); for row in (column + 3)..dimension { - let candidate = data[row * dimension + column].abs(); + let candidate = data[row * dimension + column].magnitude(); if candidate > best { best = candidate; pivot_row = row; } } - if best > data[(column + 1) * dimension + column].abs() { + if best > data[(column + 1) * dimension + column].magnitude() { // Congruence swap of rows then columns column+1 <-> pivot_row. for index in 0..dimension { data.swap((column + 1) * dimension + index, pivot_row * dimension + index); @@ -50,7 +104,7 @@ fn pfaffian_one(matrix: Array2) -> T { sign = -sign; } let pivot = data[(column + 1) * dimension + column]; - if pivot.abs() < epsilon { + if pivot.magnitude() < PIVOT_EPSILON { return T::zero(); } // Rank-2 skew Schur-complement update on the trailing block, read from originals. @@ -83,7 +137,20 @@ fn pfaffian_one(matrix: Array2) -> T { /// /// The batch elements are independent, so they are mapped over rayon threads; the per-matrix /// Parlett-Reid elimination itself stays sequential. The caller releases the GIL around this. -fn signed_pfaffian_owned(matrices: Vec>) -> Vec { +fn signed_pfaffian_owned(matrices: Vec>) -> Vec +where + T: Copy + + Zero + + One + + Magnitude + + Neg + + Add + + Sub + + Mul + + Div + + Send + + Sync, +{ if matrices.len() >= PARALLEL_BATCH_THRESHOLD { matrices.into_par_iter().map(pfaffian_one).collect() } else { @@ -92,7 +159,7 @@ fn signed_pfaffian_owned(matrices: Vec>) -> Ve } /// Copy each ``(n, n)`` slice of a ``(batch, n, n)`` view into an owned matrix. -fn owned_matrices(matrix: &PyReadonlyArray3<'_, T>) -> Vec> { +fn owned_matrices(matrix: &PyReadonlyArray3<'_, T>) -> Vec> { let view = matrix.as_array(); let batch = view.shape()[0]; (0..batch).map(|index| view.index_axis(Axis(0), index).to_owned()).collect() @@ -114,16 +181,54 @@ fn signed_pfaffian_f32<'py>(py: Python<'py>, matrix: PyReadonlyArray3<'py, f32>) Array1::from(results).into_pyarray(py) } +/// Signed Pfaffian of a batch of ``float16`` skew-symmetric matrices, shape ``(batch, n, n)``. +/// +/// The elimination runs entirely in half precision, so every intermediate is rounded to ``float16``. +/// This is the least accurate kernel and can overflow the narrow ``float16`` range. +#[pyfunction] +fn signed_pfaffian_f16<'py>(py: Python<'py>, matrix: PyReadonlyArray3<'py, f16>) -> Bound<'py, PyArray1> { + let matrices = owned_matrices(&matrix); + let results = py.allow_threads(|| signed_pfaffian_owned(matrices)); + Array1::from(results).into_pyarray(py) +} + +/// Signed Pfaffian of a batch of ``complex128`` skew-symmetric matrices, shape ``(batch, n, n)``. +#[pyfunction] +fn signed_pfaffian_c128<'py>( + py: Python<'py>, + matrix: PyReadonlyArray3<'py, Complex64>, +) -> Bound<'py, PyArray1> { + let matrices = owned_matrices(&matrix); + let results = py.allow_threads(|| signed_pfaffian_owned(matrices)); + Array1::from(results).into_pyarray(py) +} + +/// Signed Pfaffian of a batch of ``complex64`` skew-symmetric matrices, shape ``(batch, n, n)``. +#[pyfunction] +fn signed_pfaffian_c64<'py>( + py: Python<'py>, + matrix: PyReadonlyArray3<'py, Complex32>, +) -> Bound<'py, PyArray1> { + let matrices = owned_matrices(&matrix); + let results = py.allow_threads(|| signed_pfaffian_owned(matrices)); + Array1::from(results).into_pyarray(py) +} + #[pymodule] fn _rust(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add_function(wrap_pyfunction!(signed_pfaffian_f64, module)?)?; module.add_function(wrap_pyfunction!(signed_pfaffian_f32, module)?)?; + module.add_function(wrap_pyfunction!(signed_pfaffian_f16, module)?)?; + module.add_function(wrap_pyfunction!(signed_pfaffian_c128, module)?)?; + module.add_function(wrap_pyfunction!(signed_pfaffian_c64, module)?)?; Ok(()) } #[cfg(test)] mod tests { use super::pfaffian_one; + use half::f16; + use num_complex::Complex; use numpy::ndarray::array; #[test] @@ -161,4 +266,63 @@ mod tests { let empty = numpy::ndarray::Array2::::zeros((0, 0)); assert_eq!(pfaffian_one(empty), 1.0); } + + #[test] + fn f16_four_by_four_matches_formula_within_half_precision() { + // Native half-precision elimination; pf = a*f - b*e + c*d, checked at f16 tolerance. + let value = |x: f64| f16::from_f64(x); + let (a, b, c, d, e, f) = (0.5, 0.25, 0.75, 0.125, 0.375, 0.625); + let zero = f16::from_f64(0.0); + let m = array![ + [zero, value(a), value(b), value(c)], + [-value(a), zero, value(d), value(e)], + [-value(b), -value(d), zero, value(f)], + [-value(c), -value(e), -value(f), zero] + ]; + let expected = a * f - b * e + c * d; + assert!((pfaffian_one(m).to_f64() - expected).abs() < 1e-2); + } + + #[test] + fn complex_two_by_two_is_the_offdiagonal() { + // pf([[0, z], [-z, 0]]) = z for complex z. + let z = Complex::new(1.0_f64, 2.0); + let m = array![[Complex::new(0.0, 0.0), z], [-z, Complex::new(0.0, 0.0)]]; + let result = pfaffian_one(m); + assert!((result - z).norm() < 1e-12); + } + + #[test] + fn complex_four_by_four_matches_pfaffian_formula() { + // pf = a*f - b*e + c*d holds over the complex field as well. + let a = Complex::new(1.0_f64, -1.0); + let b = Complex::new(0.5, 2.0); + let c = Complex::new(-1.5, 0.25); + let d = Complex::new(2.0, 1.0); + let e = Complex::new(-0.5, -0.5); + let f = Complex::new(3.0, -2.0); + let zero = Complex::new(0.0, 0.0); + let m = array![ + [zero, a, b, c], + [-a, zero, d, e], + [-b, -d, zero, f], + [-c, -e, -f, zero] + ]; + let expected = a * f - b * e + c * d; + assert!((pfaffian_one(m) - expected).norm() < 1e-9); + } + + #[test] + fn complex_singular_is_zero() { + // A complex skew matrix with a zero pivot column has Pfaffian 0. + let zero = Complex::new(0.0_f64, 0.0); + let z = Complex::new(1.0, 1.0); + let m = array![ + [zero, zero, zero, zero], + [zero, zero, zero, zero], + [zero, zero, zero, z], + [zero, zero, -z, zero] + ]; + assert!(pfaffian_one(m).norm() < 1e-12); + } } diff --git a/src/torch_pfaffian/__init__.py b/src/torch_pfaffian/__init__.py index a053b4c..9182bf5 100644 --- a/src/torch_pfaffian/__init__.py +++ b/src/torch_pfaffian/__init__.py @@ -54,8 +54,11 @@ def pfaffian(matrix: torch.Tensor, *, sign: bool = True, check_input: bool = Fal ``sign=False``, no grad ``PfaffianDet`` cheapest: ``sqrt(|det|)`` only =================================== =========================== ========================================= - The Rust kernel runs on CPU, so a non-CPU (e.g. CUDA) input is routed to ``PfaffianParlettReid``, - which runs natively on the input device and avoids a host round-trip. + The Rust kernel runs on CPU (in real and complex precisions), so a non-CPU (e.g. CUDA) input is + routed to ``PfaffianParlettReid``, which runs natively on the input device and avoids a host + round-trip. Both strategies compute the correct complex signed Pfaffian without discarding the + imaginary part and are differentiable end-to-end for real and complex inputs (the complex backward + follows PyTorch's Wirtinger convention). The Pfaffian is only defined for skew-symmetric matrices; the strategies assume this and do not check it. Pass ``check_input=True`` to validate the assumption. For large matrices the Pfaffian diff --git a/src/torch_pfaffian/strategies/pfaffian_rust_parlett_reid.py b/src/torch_pfaffian/strategies/pfaffian_rust_parlett_reid.py index 426e652..82d5324 100644 --- a/src/torch_pfaffian/strategies/pfaffian_rust_parlett_reid.py +++ b/src/torch_pfaffian/strategies/pfaffian_rust_parlett_reid.py @@ -11,18 +11,33 @@ class RustPfaffianParlettReid(PfaffianStrategy): Compute the signed Pfaffian with the Rust Parlett-Reid kernel. The forward moves the input to a contiguous CPU array and calls the compiled - ``torch_pfaffian._rust`` kernel at a precision chosen from the input dtype: ``float32`` inputs - use the single-precision kernel and every other floating dtype uses the double-precision kernel. + ``torch_pfaffian._rust`` kernel for the input dtype, each running natively at that precision: + ``float16``, ``float32``/``complex64`` and ``float64``/``complex128`` map to the half-, single- + and double-precision kernels. Half precision is the caller's explicit choice and carries its risk: + the elimination runs entirely in ``float16``, which is the least accurate kernel and can overflow + the narrow ``float16`` range to ``inf`` for larger or larger-scaled matrices (use ``float32`` for a + measurably more accurate result). Complex inputs are computed natively over the complex field (no + imaginary part is discarded), giving the correct complex signed Pfaffian. Any other dtype raises + :class:`TypeError`. The result is cast back to the input dtype and device. The backward is the same as - :class:`PfaffianParlettReid` (the Pfaffian adjugate ``d pf(A) / d A = (1 / 2) pf(A) (A^{-1})^T``, - exact for invertible and singular inputs), computed in PyTorch. CUDA inputs are evaluated on CPU - for the forward; the backward runs on the input device. + :class:`PfaffianParlettReid` (the Pfaffian adjugate + ``d pf(A) / d A = (1 / 2) pf(A) (A^{-1})^T``, exact for invertible and singular inputs and + conjugated for complex autograd), computed in PyTorch. CUDA inputs are evaluated on CPU for the + forward; the backward runs on the input device. The input is a skew-symmetric matrix of shape ``(..., n, n)``. """ NAME = "RustPfaffianParlettReid" + _KERNEL_BY_DTYPE = { + torch.float16: ("signed_pfaffian_f16", torch.float16), + torch.float32: ("signed_pfaffian_f32", torch.float32), + torch.float64: ("signed_pfaffian_f64", torch.float64), + torch.complex64: ("signed_pfaffian_c64", torch.complex64), + torch.complex128: ("signed_pfaffian_c128", torch.complex128), + } + @staticmethod def forward(matrix: torch.Tensor) -> torch.Tensor: dimension = matrix.shape[-1] @@ -30,13 +45,17 @@ def forward(matrix: torch.Tensor) -> torch.Tensor: return torch.zeros(matrix.shape[:-2], dtype=matrix.dtype, device=matrix.device) if dimension == 0: return torch.ones(matrix.shape[:-2], dtype=matrix.dtype, device=matrix.device) + if matrix.dtype not in RustPfaffianParlettReid._KERNEL_BY_DTYPE: + supported = ", ".join(str(dtype) for dtype in RustPfaffianParlettReid._KERNEL_BY_DTYPE) + raise TypeError( + f"RustPfaffianParlettReid has no Rust kernel for dtype {matrix.dtype}; supported dtypes " + f"are {supported}. Cast the matrix to a supported dtype, or use the pure-PyTorch " + "PfaffianParlettReid strategy which handles any dtype. To request a Rust kernel for this " + "dtype, please open an issue at https://github.com/MatchCake/TorchPfaffian/issues." + ) flat = matrix.reshape(-1, dimension, dimension) # (batch, n, n) - if matrix.dtype == torch.float32: - working_dtype = torch.float32 - kernel = _rust.signed_pfaffian_f32 - else: - working_dtype = torch.float64 - kernel = _rust.signed_pfaffian_f64 + kernel_name, working_dtype = RustPfaffianParlettReid._KERNEL_BY_DTYPE[matrix.dtype] + kernel = getattr(_rust, kernel_name) array = flat.detach().to(working_dtype).cpu().contiguous().numpy() result = kernel(array) # (batch,) in working_dtype pfaffian = torch.from_numpy(result).to(dtype=matrix.dtype, device=matrix.device) diff --git a/src/torch_pfaffian/strategies/strategy.py b/src/torch_pfaffian/strategies/strategy.py index 171cb8e..c457960 100644 --- a/src/torch_pfaffian/strategies/strategy.py +++ b/src/torch_pfaffian/strategies/strategy.py @@ -42,6 +42,11 @@ def pfaffian_grad_matrix( ``pf == 0`` elements (whose inverse is discarded anyway) are replaced by the identity before the batched inverse so the call stays well-posed. + The Pfaffian is holomorphic in the entries of ``A``, so for complex inputs the backward returns + the conjugate of the analytic derivative, ``conj(d pf / d A) * grad_output``, which is PyTorch's + Wirtinger convention for complex autograd (``z.grad = d L / d conj(z)``). For real inputs the + conjugation is a no-op, so real gradients are unchanged. + :param matrix: The saved input matrix of shape ``(..., n, n)``. :param pfaffian: The saved forward Pfaffian of shape ``(...,)``. :param grad_output: Gradient of the output with respect to the loss, of shape ``(...,)``. @@ -66,7 +71,7 @@ def pfaffian_grad_matrix( singular_adjugate = cls._pfaffian_adjugate(flat_matrix.index_select(0, singular_index)) flat_adjugate = flat_adjugate.index_copy(0, singular_index, singular_adjugate.to(flat_adjugate.dtype)) adjugate = flat_adjugate.reshape_as(matrix) - return torch.einsum("...,...ij->...ji", 0.5 * grad_output, adjugate) + return torch.einsum("...,...ij->...ji", 0.5 * grad_output, adjugate.conj()) @classmethod def _pfaffian_adjugate(cls, matrices: torch.Tensor) -> torch.Tensor: diff --git a/tests/test_strategies/test_pfaffian_rust_parlett_reid.py b/tests/test_strategies/test_pfaffian_rust_parlett_reid.py index 61ad21f..093c0b4 100644 --- a/tests/test_strategies/test_pfaffian_rust_parlett_reid.py +++ b/tests/test_strategies/test_pfaffian_rust_parlett_reid.py @@ -1,3 +1,5 @@ +import warnings + import numpy as np import pytest import torch @@ -37,6 +39,28 @@ def _random_skew(dimension: int, rng: np.random.Generator) -> torch.Tensor: return torch.tensor(skew - skew.T) +def _rand_skew_complex(dimension: int, rng: np.random.Generator) -> np.ndarray: + entries = rng.normal(size=(dimension, dimension)) + 1j * rng.normal(size=(dimension, dimension)) + return entries - entries.T + + +def _cofactor_pfaffian(matrix: np.ndarray) -> complex: + # Exact recursive cofactor expansion; ground-truth oracle for complex skew-symmetric inputs. + dimension = matrix.shape[0] + if dimension == 0: + return 1 + 0j + if dimension % 2: + return 0j + if dimension == 2: + return matrix[0, 1] + total = 0j + rest = list(range(1, dimension)) + for position, column in enumerate(rest): + sub = [index for index in rest if index != column] + total += (-1) ** position * matrix[0, column] * _cofactor_pfaffian(matrix[np.ix_(sub, sub)]) + return total + + def _skew_from_parameters(parameters: torch.Tensor, dimension: int) -> torch.Tensor: upper = torch.zeros(dimension, dimension, dtype=parameters.dtype) indices = torch.triu_indices(dimension, dimension, offset=1) @@ -165,6 +189,68 @@ def test_backward_gradcheck_at_singular_point(self): rtol=RTOL_APPROX_COMPARISON, ) + @pytest.mark.parametrize("dimension", [2, 4, 6, 8]) + @pytest.mark.parametrize("dtype", [torch.complex64, torch.complex128]) + def test_forward_complex_matches_cofactor_oracle(self, dimension, dtype): + rng = np.random.default_rng(TEST_SEED + dimension) + skew = _rand_skew_complex(dimension, rng) + matrix = torch.tensor(skew, dtype=dtype) + result = RustPfaffianParlettReid.apply(matrix) + assert result.dtype == dtype + atol = ATOL_SCALAR_COMPARISON if dtype == torch.complex128 else 1e-4 + torch.testing.assert_close(result, torch.tensor(_cofactor_pfaffian(skew), dtype=dtype), atol=atol, rtol=atol) + + def test_forward_complex_matches_python_parlett_reid(self): + rng = np.random.default_rng(TEST_SEED) + matrix = torch.tensor(_rand_skew_complex(6, rng), dtype=torch.complex128) + torch.testing.assert_close( + RustPfaffianParlettReid.apply(matrix), + PfaffianParlettReid.apply(matrix), + atol=ATOL_SCALAR_COMPARISON, + rtol=RTOL_SCALAR_COMPARISON, + ) + + def test_forward_complex_emits_no_imaginary_discard_warning(self): + rng = np.random.default_rng(TEST_SEED) + matrix = torch.tensor(_rand_skew_complex(6, rng), dtype=torch.complex128) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + RustPfaffianParlettReid.apply(matrix) + assert not any("imaginary part" in str(warning.message) for warning in caught) + + @pytest.mark.parametrize("dimension", _GRADCHECK_DIMENSIONS) + def test_backward_complex_passes_gradcheck_on_skew_parameterization(self, dimension): + count = dimension * (dimension - 1) // 2 + rng = np.random.default_rng(TEST_SEED + dimension) + values = rng.normal(size=count) + 1j * rng.normal(size=count) + parameters = torch.tensor(values, dtype=torch.complex128, requires_grad=True) + assert gradcheck( + lambda free: RustPfaffianParlettReid.apply(_skew_from_parameters(free, dimension)), + (parameters,), + eps=1e-6, + atol=ATOL_APPROX_COMPARISON, + rtol=RTOL_APPROX_COMPARISON, + ) + + def test_forward_float16_runs_natively_in_half_precision(self): + # float16 is computed by the native half-precision kernel (the caller opts into half-precision + # risk); the result agrees with the double-precision Pfaffian of the same f16 values at f16 + # tolerance and preserves the float16 dtype. + half = _random_skew(4, _RNG).to(torch.float16) + result = RustPfaffianParlettReid.apply(half) + assert result.dtype == torch.float16 + reference = RustPfaffianParlettReid.apply(half.to(torch.float64)) + torch.testing.assert_close( + result.to(torch.float64), reference, atol=ATOL_APPROX_COMPARISON, rtol=RTOL_APPROX_COMPARISON + ) + + def test_forward_unsupported_dtype_raises_with_guidance(self): + # Unsupported dtypes must raise rather than be silently downcast; the message points at the + # pure-PyTorch fallback and the issue tracker. + matrix = torch.tensor([[0.0, 1.0], [-1.0, 0.0]]).to(torch.bfloat16) + with pytest.raises(TypeError, match="no Rust kernel for dtype.*PfaffianParlettReid.*issues"): + RustPfaffianParlettReid.apply(matrix) + def test_backward_returns_none_when_input_does_not_require_grad(self): matrix = _random_skew(4, _RNG) pfaffian = RustPfaffianParlettReid.apply(matrix) diff --git a/tests/test_torch_pfaffian.py b/tests/test_torch_pfaffian.py index cc98b7e..5eaed78 100644 --- a/tests/test_torch_pfaffian.py +++ b/tests/test_torch_pfaffian.py @@ -1,13 +1,43 @@ +import warnings from unittest import mock +import numpy as np import pytest import torch import torch_pfaffian +from tests.configs import ( + ATOL_SCALAR_COMPARISON, + N_RANDOM_TESTS_PER_CASE, + RTOL_SCALAR_COMPARISON, + TEST_SEED, +) from torch_pfaffian import get_pfaffian_function, pfaffian, pfaffian_strategy_map from torch_pfaffian.strategies import PfaffianDet, PfaffianFDBPf, PfaffianParlettReid +def _cofactor_pfaffian(matrix: np.ndarray) -> complex: + # Exact recursive cofactor expansion; ground-truth oracle for complex skew-symmetric inputs. + dimension = matrix.shape[0] + if dimension == 0: + return 1 + 0j + if dimension % 2: + return 0j + if dimension == 2: + return matrix[0, 1] + total = 0j + rest = list(range(1, dimension)) + for position, column in enumerate(rest): + sub = [index for index in rest if index != column] + total += (-1) ** position * matrix[0, column] * _cofactor_pfaffian(matrix[np.ix_(sub, sub)]) + return total + + +def _rand_skew_complex(dimension: int, rng: np.random.Generator) -> np.ndarray: + entries = rng.normal(size=(dimension, dimension)) + 1j * rng.normal(size=(dimension, dimension)) + return entries - entries.T + + class TestTorchPfaffian: def test_rust_parlett_reid_registered_when_available(self): pytest.importorskip("torch_pfaffian._rust") @@ -127,6 +157,107 @@ def test_pfaffian_check_input_off_by_default_allows_non_skew(self): non_skew = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float64) pfaffian(non_skew) # no validation by default, so no error is raised + @pytest.mark.parametrize("dimension", [2, 4, 6, 8, 10, 12]) + @pytest.mark.parametrize("dtype", [torch.complex64, torch.complex128]) + def test_pfaffian_sign_true_matches_complex_oracle(self, dimension, dtype): + rng = np.random.default_rng(TEST_SEED + dimension) + for _ in range(N_RANDOM_TESTS_PER_CASE): + skew = _rand_skew_complex(dimension, rng) + expected = _cofactor_pfaffian(skew) + matrix = torch.tensor(skew, dtype=dtype) + result = pfaffian(matrix, sign=True) + assert result.dtype == dtype + atol = ATOL_SCALAR_COMPARISON if dtype == torch.complex128 else 1e-4 + torch.testing.assert_close( + result, + torch.tensor(expected, dtype=dtype), + atol=atol, + rtol=atol, + ) + assert abs(expected.imag) > 1e-6 # the oracle carries a genuine imaginary part + + def test_pfaffian_sign_true_complex_squares_to_det(self): + rng = np.random.default_rng(TEST_SEED) + matrix = torch.tensor(_rand_skew_complex(6, rng), dtype=torch.complex128) + result = pfaffian(matrix, sign=True) + torch.testing.assert_close( + result**2, torch.linalg.det(matrix), atol=ATOL_SCALAR_COMPARISON, rtol=RTOL_SCALAR_COMPARISON + ) + + def test_pfaffian_sign_true_complex_supports_batched_shapes(self): + rng = np.random.default_rng(TEST_SEED) + batch = np.stack([_rand_skew_complex(4, rng) for _ in range(6)]).reshape(2, 3, 4, 4) + matrix = torch.tensor(batch, dtype=torch.complex128) + result = pfaffian(matrix, sign=True) + assert result.shape == (2, 3) + flat = matrix.reshape(-1, 4, 4) + expected = torch.stack([pfaffian(flat[index], sign=True) for index in range(flat.shape[0])]).reshape(2, 3) + torch.testing.assert_close(result, expected, atol=ATOL_SCALAR_COMPARISON, rtol=RTOL_SCALAR_COMPARISON) + + def test_pfaffian_sign_true_complex_routes_to_rust_on_cpu(self): + # The Rust kernel now handles complex natively, so complex CPU inputs take the fast Rust path. + pytest.importorskip("torch_pfaffian._rust") + rng = np.random.default_rng(TEST_SEED) + matrix = torch.tensor(_rand_skew_complex(4, rng), dtype=torch.complex128) + with mock.patch.object(torch_pfaffian, "RustPfaffianParlettReid") as fake_rust: + fake_rust.apply.return_value = torch.zeros((), dtype=torch.complex128) + pfaffian(matrix, sign=True) + fake_rust.apply.assert_called_once() + + def test_pfaffian_sign_true_complex_emits_no_imaginary_discard_warning(self): + rng = np.random.default_rng(TEST_SEED) + matrix = torch.tensor(_rand_skew_complex(6, rng), dtype=torch.complex128) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + pfaffian(matrix, sign=True) + assert not any("imaginary part" in str(warning.message) for warning in caught) + + def test_pfaffian_sign_true_complex_backward_passes_gradcheck(self): + dimension = 4 + count = dimension * (dimension - 1) // 2 + rng = np.random.default_rng(TEST_SEED) + values = rng.normal(size=count) + 1j * rng.normal(size=count) + parameters = torch.tensor(values, dtype=torch.complex128, requires_grad=True) + + def from_parameters(free): + upper = torch.zeros(dimension, dimension, dtype=free.dtype) + indices = torch.triu_indices(dimension, dimension, offset=1) + upper = upper.index_put((indices[0], indices[1]), free) + return upper - upper.transpose(-1, -2) + + assert torch.autograd.gradcheck( + lambda free: pfaffian(from_parameters(free), sign=True), (parameters,), eps=1e-6, atol=1e-4, rtol=1e-4 + ) + + def test_pfaffian_sign_true_complex_odd_dimension_is_zero(self): + rng = np.random.default_rng(TEST_SEED) + matrix = torch.tensor(_rand_skew_complex(5, rng), dtype=torch.complex128) + torch.testing.assert_close( + pfaffian(matrix, sign=True), + torch.zeros((), dtype=torch.complex128), + atol=ATOL_SCALAR_COMPARISON, + rtol=RTOL_SCALAR_COMPARISON, + ) + + def test_pfaffian_sign_true_complex_empty_is_one(self): + matrix = torch.zeros((0, 0), dtype=torch.complex128) + torch.testing.assert_close( + pfaffian(matrix, sign=True), + torch.ones((), dtype=torch.complex128), + atol=ATOL_SCALAR_COMPARISON, + rtol=RTOL_SCALAR_COMPARISON, + ) + + def test_pfaffian_sign_true_complex_singular_is_zero_without_nan(self): + matrix = torch.zeros((4, 4), dtype=torch.complex128) + matrix[2, 3] = 1.0 + 1.0j + matrix[3, 2] = -(1.0 + 1.0j) + result = pfaffian(matrix, sign=True) + assert torch.isfinite(result.real) and torch.isfinite(result.imag) + torch.testing.assert_close( + result, torch.zeros((), dtype=torch.complex128), atol=ATOL_SCALAR_COMPARISON, rtol=RTOL_SCALAR_COMPARISON + ) + def test_pfaffian_warns_when_result_overflows(self): # A 4x4 block-antidiagonal with huge entries makes the Pfaffian overflow to inf. block = torch.tensor([[1e200, 0.0], [0.0, 1e200]], dtype=torch.float64)