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/3] 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/3] 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/3] 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(