TorchPfaffian

torch_pfaffian package

Subpackages

Submodules

torch_pfaffian.utils module

torch_pfaffian.utils.get_all_subclasses(cls)

Module contents

TorchPfaffian computes batched, differentiable Pfaffians of skew-symmetric matrices in PyTorch, on CPU and GPU.

torch_pfaffian.get_pfaffian_function(name: str = 'PfaffianFDBPf') → Callable[[Tensor], Tensor]
torch_pfaffian.pfaffian(matrix: Tensor, *, sign: bool = True, check_input: bool = False, check_finite: bool = True, epsilon: float | None = None, mask_roundoff_gradient: bool = False) → Tensor

Compute the Pfaffian of a skew-symmetric matrix, choosing the strategy from the input.

The matrix has shape (..., 2n, 2n) and the result has shape (...,), sharing the backend, dtype, and device of matrix.

Strategy selection:

Condition

Strategy

Reason

mask_roundoff_gradient=True

PfaffianMaskedMagnitude

zero gradient at round-off zeros

m <= AUTO_SMALL_MAX

PfaffianSmall

exact unrolled kernel (all dtypes)

block form, large input (see below)

PfaffianBlockDet

one n x n determinant

sign=True, Rust built, CPU input

RustPfaffianParlettReid

fastest signed path (native Rust)

sign=True, otherwise

PfaffianParlettReid

GPU-native and pure-Python fallback

sign=False, epsilon given

magnitude_pfaffian

log-domain magnitude with floor

sign=False, grad needed

PfaffianFDBPf

magnitude only; robust backward

sign=False, no grad

PfaffianDet

cheapest: sqrt(|det|) only

Matrices with last dimension m <= AUTO_SMALL_MAX take the exact unrolled PfaffianSmall path first, regardless of sign and device (|Pf| for sign=False); it is exact for every dtype and differentiable end-to-end. For larger matrices 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. All signed 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).

Inputs with m > AUTO_SMALL_MAX and at least AUTO_BLOCK_DET_MIN_NUMEL entries in total are tested for the block form [[0, B], [-B^T, 0]] (both diagonal blocks exactly zero for every batch element). When it holds, PfaffianBlockDet computes pf = (-1)^{n(n-1)/2} det(B) from the n x n block, for sign=True and for sign=False without epsilon (where the magnitude is floored at sqrt(PfaffianStrategy.EPSILON) like the determinant paths). The test rejects a general input after reading one entry per matrix, and the gradient matches the general strategies.

The supported dtypes are float16, bfloat16, float32, float64, complex32, complex64 and complex128. PyTorch has no LU factorization in float16, bfloat16 and complex32, so these inputs are computed in float32 (complex64 for complex32), with a UserWarning, and the result is cast back: at every size for complex32, and for float16 and bfloat16 when m > AUTO_SMALL_MAX, with mask_roundoff_gradient=True, or when the unrolled kernel’s result is not finite. The gradient flows back through the cast. A result that is finite after the upcast but exceeds the range of the input dtype becomes inf, with a RuntimeWarning when check_finite=True. The round-off criterion of mask_roundoff_gradient then uses the precision of the upcast dtype, so a Pfaffian that is zero only up to the round-off of an earlier half-precision computation is not masked.

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. The chosen strategy is the fastest correct one for the input, so calling pfaffian is all most callers need. For inputs with a wide dynamic range the fast linear-domain path can overflow to inf even when the true Pfaffian is finite; with check_finite=True (default) such results are transparently recomputed in the log domain and recovered, so the caller gets the correct value without reaching for slog_pfaffian(). Only a magnitude that genuinely exceeds the input dtype’s floating range remains non-finite, and only then is a RuntimeWarning emitted. One extreme case is not auto-recovered: for m > AUTO_SMALL_MAX the linear signed path declares a pivot magnitude below PfaffianStrategy.EPSILON (1e-30) an exact zero, so a matrix whose Pfaffian hinges on such a sub-1e-30 pivot is reported as 0. slog_pfaffian(), which uses an exact-zero pivot test, computes these extreme inputs correctly (as does the unrolled m <= AUTO_SMALL_MAX path, which has no pivot floor).

Parameters:
  • matrix – Skew-symmetric matrix of shape (..., 2n, 2n).

  • sign – When True (default) return the signed Pfaffian, otherwise its magnitude.

  • check_input – When True, validate that matrix is square in its last two dimensions and skew-symmetric (A == -A^T) before computing, raising ValueError otherwise. Off by default (False) so trusted inputs pay nothing; the check is an O(n^2) comparison, cheap relative to the O(n^3) Pfaffian.

  • check_finite – When True (default), verify the result is finite; if the fast path overflowed, recompute in the log domain and recover the value, warning only if it still exceeds the dtype’s range. This costs one device-to-host synchronization; pass False to skip both the check and the recovery in hot loops (the raw fast-path result, possibly inf, is returned).

  • epsilon – Only affects the magnitude path (sign=False). When a float is given, the magnitude is computed in the log domain and floored at sqrt(epsilon); elements below the floor get a zero gradient. When None (default), the routing is unchanged and the magnitude is computed via PfaffianFDBPf / PfaffianDet (or PfaffianBlockDet for large block-form inputs) and floored at sqrt(PfaffianStrategy.EPSILON). With mask_roundoff_gradient=True it is the per-call floor of |det(A)| (None means PfaffianStrategy.EPSILON).

  • mask_roundoff_gradient – Only valid with sign=False. When True, every input size is computed by PfaffianMaskedMagnitude: the forward is sqrt(max(|det(A)|, epsilon)) and the backward gives a zero gradient to elements whose Pfaffian vanishes up to round-off (or that were clamped to the floor), instead of the O(1) gradient the unmasked backward derives from round-off. No numerical test separates such a structural zero from a genuine value at round-off level, so this is off by default (False); turn it on when exact zeros are expected, for example the parity-forbidden outcomes of a basis state. The backward costs one inverse, two 1-norms and one set of singular values per element. Elements whose determinant overflows are recovered in the log domain with the unmasked gradient; see PfaffianMaskedMagnitude for the limits of the round-off criterion.

Raises:
  • ValueError – If mask_roundoff_gradient=True with sign=True, or if check_input=True and matrix is not square and skew-symmetric.

  • TypeError – If matrix has an unsupported dtype (an integer, boolean or float8 dtype).

Returns:

The Pfaffian of the input, of shape (...,).

Return type:

torch.Tensor

torch_pfaffian.slog_pfaffian(matrix: Tensor) → tuple[torch.Tensor, torch.Tensor]

Log-domain signed Pfaffian of a batch of skew-symmetric matrices as (phase, log|Pf|).

Thin wrapper over SlogPfaffianStrategy. The reconstruction is Pf = phase * exp(log_abs). The log domain neither overflows for large matrices nor underflows for tiny Pfaffians, and the kernel never synchronizes with the host. log_abs is differentiable wherever Pf != 0; phase is non-differentiable (use pfaffian() with sign=True for the full derivative). float16 and bfloat16 inputs are computed in float32, and complex32 inputs in complex64, with a UserWarning; the result is cast back, since PyTorch has no LU factorization in those dtypes.

Parameters:

matrix – Skew-symmetric matrix of shape (..., 2n, 2n).

Returns:

A pair (phase, log_abs), each of shape (...,). phase shares the input dtype (0 when Pf = 0); log_abs is in the matching real dtype (-inf when Pf = 0).

Return type:

tuple[torch.Tensor, torch.Tensor]

Raises:

TypeError – If matrix has an unsupported dtype (an integer, boolean or float8 dtype).