Fix bwk

src/torch_pfaffian/strategies/pfaffian_block_det.py

@@ -29,31 +29,75 @@ 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.
         :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]:
-            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

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)

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)``.
@@ -39,31 +39,29 @@ 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
             valid = valid & ~zero_pivot
             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

src/torch_pfaffian/strategies/strategy.py

@@ -27,22 +27,38 @@ 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
-        if bool(singular.any()):
-            dimension = matrix.shape[-1]
+        dimension = matrix.shape[-1]
+        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]

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