torch_vn_algebra.functional¶

Source: torch_vn_algebra/functional.py. Signatures and docstrings as in the source.

Module documentation¶

Functional calculus on stacks of blocks (n, k, k), at four levels of generality.

1. Self-adjoint X: f(X) = V f(w) V^* from the Hermitian eigensolver; f any function on the real spectrum.
2. Normal X (unitaries, ...): the eigenvectors of the Hermitian matrix H + alpha K, X = H + i K, diagonalise X
   (H and K commute); f on the complex spectrum. Again only the Hermitian eigensolver is used.
3. Singular values of any X from the Hermitian eigensolver of the dilation [[0, X], [X^*, 0]], whose
   eigenvalues are +-s_i: |X|, the pseudo-inverse, the polar decomposition and the norms. Batched GPU SVD
   is limited to blocks of size <= 32; the Hermitian eigensolver is batched for every size.
4. Any X and f holomorphic near the spectrum: f(X) = (2 pi i)^{-1} oint f(z) (z - X)^{-1} dz, by the
   trapezoidal rule on a circle (batched LU solves), with an error estimate from doubling the nodes.

Gradients. Levels 1 and 2 use the Daleckii-Krein formula dF = V (L o V^* dX V) V^* with the Loewner matrix
of divided differences of f (f' for coinciding eigenvalues), which is finite at degenerate spectra; level 2
assumes f holomorphic near the spectrum (the derivative is then the Frechet derivative of the holomorphic
calculus in every direction). Second derivatives (create_graph) recompute the decomposition differentiably
(eigh for self-adjoint, eig for normal X) and are correct at non-degenerate spectra. Level 3 has explicit
backward formulas (L_ij = 1/(s_i + s_j) for |X| on the range, the constant-rank formula for the pseudo-inverse);
level 4 is differentiated through the linear solves.

cache_enabled¶

def cache_enabled() -> bool

caching¶

def caching(enabled: bool=True)

Within the block, every operator keeps its eigendecomposition after the first use (see
Operator.cache_spectrum): repeated functions of the same operator then cost matrix products only.

values_and_derivative¶

def values_and_derivative(f: Callable[[torch.Tensor], torch.Tensor], w: torch.Tensor, create_graph: bool=False) -> Tuple[torch.Tensor, torch.Tensor]

f(w) and f'(w) for an elementwise f, with one evaluation of f. For real w, f' is the real derivative;
for complex w, f' is the complex derivative of a holomorphic f (PyTorch returns conj(f') for the
gradient of Re f). The values keep their graph to parameters of f. Non-floating results (indicators)
are converted to the dtype of w; f' is 0 where f does not depend on w.

divided_differences¶

def divided_differences(w: torch.Tensor, fw: torch.Tensor, dfw: torch.Tensor) -> torch.Tensor

Loewner matrix L_ij = f[w_i, w_j] = (f(w_i) - f(w_j)) / (w_i - w_j) (f'(w_i) on ties) of the Daleckii-Krein
formula (Theorem 1.11 of the lecture notes), from the values fw = f(w) and dfw = f'(w). w may be complex.
Shape (..., k, k).

For close w_i, w_j the quotient q loses accuracy by cancellation, and the midpoint m = (f'(w_i) + f'(w_j)) / 2,
whose error is at most |w_i - w_j|^2 sup |f'''| / 12 (Proposition 1.12), may be better. The choice
(Proposition 1.13) compares the two with a bound R for the rounding error of q: m is used iff |q - m| <= 2 R,
which makes the error at most min(T, 3 R), T the error of m, without knowing f'''. Here
R = eps (d_i + d_j) / |w_i - w_j| + 4 eps |q| with d_i = |f(w_i)| + |w_i f'(w_i)|: the assumed evaluation error
of f (relative error eps of the value, at an argument with relative error eps; e.g. exp(-1j * t * w) rounds
t * w). The midpoint is also used for ties, for a non-finite quotient, and for pairs that the eigensolver does
not resolve, |w_i - w_j| <= k eps max_l |w_l| (e.g. a degenerate eigenvalue 0 computed as +-1e-17, where the
quotient is meaningless also for a discontinuous f).

hermitian_eig¶

def hermitian_eig(X: torch.Tensor, what: str='eigh') -> Tuple[torch.Tensor, torch.Tensor]

Eigenvalues (real, ascending) and eigenvectors of the self-adjoint part of a stack X.

normal_eig¶

def normal_eig(X: torch.Tensor, what: str='eigh', rtol: Optional[float]=None) -> Tuple[torch.Tensor, torch.Tensor]

Eigenvalues (complex) and a unitary eigenbasis of a stack X of normal matrices, from the Hermitian
eigensolver: X = H + i K with commuting self-adjoint H, K, and the eigenvectors of H + alpha K diagonalise
both when alpha separates the eigenvalues (true for all but finitely many alpha). Eigenvectors of nearly
tied values of Re + alpha Im are refined with a second coefficient; several irrational coefficients are tried
and the best result is kept per matrix. Raises ValueError if some X is not normal to working precision
(relative off-diagonal residual above rtol, default 100 k eps).

spectral_apply¶

def spectral_apply(X: torch.Tensor, f: Callable[[torch.Tensor], torch.Tensor], kind: str='hermitian', decomposition: Optional[Tuple[torch.Tensor, torch.Tensor]]=None, what: str='eigh') -> torch.Tensor

f(X) for a stack X (n, k, k) of self-adjoint (kind='hermitian') or normal (kind='normal') matrices and an
elementwise f on the eigenvalues (n, k). `decomposition` = (w, V) may be passed from a cache. For
self-adjoint X the function is evaluated at the self-adjoint part of X.

use_dilation¶

def use_dilation(X: torch.Tensor, method: str='auto') -> bool

'auto': the dilation on GPUs for blocks larger than 32 (where batched SVD falls back to a loop).

singular_values¶

def singular_values(X: torch.Tensor, method: str='auto') -> torch.Tensor

Singular values (n, k), descending, of a stack of square blocks.

svd¶

def svd(X: torch.Tensor, method: str='auto') -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]

X = U diag(s) V^*, without gradients (callers attach their own backward). With the dilation, the columns
of U and V belonging to zero singular values are not determined (they never enter the results below).

abs_¶

def abs_(X: torch.Tensor, method: str='auto') -> torch.Tensor

pinv¶

def pinv(X: torch.Tensor, tol: Optional[float]=None, method: str='auto') -> torch.Tensor

polar¶

def polar(X: torch.Tensor, method: str='auto') -> Tuple[torch.Tensor, torch.Tensor]

X = W |X| with W = U V^* (unitary for invertible X; on the kernel of a singular X the factor W is the
partial isometry from the SVD) and |X| from _Abs (differentiable). W has no gradient. Two Newton-Schulz steps
W <- (3W - W W^* W)/2 restore orthonormality lost to rounding (the dilation mixes the eigenvectors of +-s
for singular values small compared with the largest).

default_contour¶

def default_contour(X: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]

Circle around the spectrum: center c = Tr(X)/k and radius 1.25 max|lambda_i - c| + 0.5 ||X - c||_F / sqrt(k)
(the second term keeps the circle away from the eigenvalues of non-normal blocks, whose resolvent is large near
the spectrum, and is positive for X = c + nilpotent). Eigenvalues from torch.linalg.eigvals, without gradients.

holomorphic¶

def holomorphic(X: torch.Tensor, f: Callable[[torch.Tensor], torch.Tensor], center=None, radius=None, nodes: int=32, max_nodes: int=2048, rtol: Optional[float]=None, memory_bytes: int=2 ** 28, atol: float=0.0) -> Tuple[torch.Tensor, float]

f(X) = (2 pi i)^{-1} oint f(z) (z - X)^{-1} dz on the circle |z - center| = radius, by the trapezoidal
rule with N nodes, N doubled until two successive results agree to rtol relative to their size (the nodes of N
are reused for 2N). Default contour: default_contour(X), a circle just around the spectrum; f must be
holomorphic on the closed disk and accept complex tensors elementwise. center and radius may be numbers,
tensors (n,) or callables X -> (n,). Raises ValueError if the estimated error stays above sqrt(rtol) (a
warning between rtol and sqrt(rtol)). The error is relative to max(|f(X)|, atol): when f(X) may vanish (sin at
X = 0, z^3 at a nilpotent block) pass atol > 0, a scale for f near the spectrum (e.g. 1.0); the absolute
accuracy accepted where |f(X)| < atol is then rtol * atol. Returns (f(X), error estimate).