torchzero 0.3.14__py3-none-any.whl → 0.4.0__py3-none-any.whl

torchzero/linalg/__init__.py

@@ -0,0 +1,10 @@
+from . import linear_operator
+
+from .matrix_power import (
+    matrix_power_eigh,
+    matrix_power_svd,
+)
+from .orthogonalize import zeropower_via_eigh, zeropower_via_newtonschulz5, zeropower_via_svd, orthogonalize
+from .qr import qr_householder
+from .solve import cg, nystrom_sketch_and_solve, nystrom_pcg
+from .eigh import nystrom_approximation

torchzero/linalg/eigh.py

@@ -0,0 +1,34 @@
+from collections.abc import Callable
+import torch
+from .linalg_utils import mm
+
+
+
+# https://arxiv.org/pdf/2110.02820
+def nystrom_approximation(
+    A_mv: Callable[[torch.Tensor], torch.Tensor] | None,
+    A_mm: Callable[[torch.Tensor], torch.Tensor] | None,
+    ndim: int,
+    rank: int,
+    device,
+    dtype = torch.float32,
+    generator = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Computes Nyström approximation to positive-semidefinite A factored as Q L Q^T (truncatd eigenvalue decomp),
+    returns ``(L, Q)``.
+
+    A is ``(m,m)``, then Q is ``(m, rank)``; L is a ``(rank, )`` vector - diagonal of ``(rank, rank)``"""
+    # basis
+    O = torch.randn((ndim, rank), device=device, dtype=dtype, generator=generator) # Gaussian test matrix
+    O, _ = torch.linalg.qr(O) # Thin QR decomposition # pylint:disable=not-callable
+
+    # Y = AΩ
+    AO = mm(A_mv=A_mv, A_mm=A_mm, X=O)
+
+    v = torch.finfo(dtype).eps * torch.linalg.matrix_norm(AO, ord='fro') # Compute shift # pylint:disable=not-callable
+    Yv = AO + v*O # Shift for stability
+    C = torch.linalg.cholesky_ex(O.mT @ Yv)[0] # pylint:disable=not-callable
+    B = torch.linalg.solve_triangular(C, Yv.mT, upper=False, unitriangular=False).mT # pylint:disable=not-callable
+    Q, S, _ = torch.linalg.svd(B, full_matrices=False) # pylint:disable=not-callable
+    L = (S.pow(2) - v).clip(min=0) #Remove shift, compute eigs
+    return L, Q

torchzero/linalg/linalg_utils.py

@@ -0,0 +1,14 @@
+from collections.abc import Callable
+import torch
+
+def mm(
+    A_mv: Callable[[torch.Tensor], torch.Tensor] | None,
+    A_mm: Callable[[torch.Tensor], torch.Tensor] | None,
+    X
+):
+    """matrix-matrix when either mv or mm is given"""
+    if A_mm is not None: return A_mm(X)
+    assert A_mv is not None
+    return torch.stack([A_mv(col) for col in X.unbind(-1)], -1) # rank matvecs
+
+

torchzero/{utils/linalg → linalg}/linear_operator.py

@@ -1,4 +1,6 @@
-"""simplified version of https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg.LinearOperator.html. This is used for trust regions."""
+"""This is mainly used for trust regions. In some cases certain operations are relaxed, e.g. eigenvalue shift instead of
+adding diagonal when it isn't tractable, to make it work with Levenberg-Marquadt.
+"""
 import math
 from abc import ABC, abstractmethod
 from functools import partial
@@ -7,7 +9,8 @@ from typing import cast, final
 
 import torch
 
-from ..torch_tools import tofloat, tonumpy, totensor
+from ..utils.torch_tools import tofloat, tonumpy, totensor
+from .solve import nystrom_sketch_and_solve
 
 if find_spec('scipy') is not None:
     from scipy.sparse.linalg import LinearOperator as _ScipyLinearOperator
@@ -15,7 +18,6 @@ else:
     _ScipyLinearOperator = None
 
 class LinearOperator(ABC):
-    """this is used for trust region"""
     device: torch.types.Device
     dtype: torch.dtype | None
 
@@ -25,18 +27,24 @@ class LinearOperator(ABC):
     def rmatvec(self, x: torch.Tensor) -> torch.Tensor:
         raise NotImplementedError(f"{self.__class__.__name__} doesn't implement rmatvec")
 
-    def matmat(self, x: torch.Tensor) -> "LinearOperator":
-        raise NotImplementedError(f"{self.__class__.__name__} doesn't implement matmul")
+    def matmat(self, X: torch.Tensor) -> "LinearOperator":
+        raise NotImplementedError(f"{self.__class__.__name__} doesn't implement matmat")
+
+    def rmatmat(self, X: torch.Tensor) -> "LinearOperator":
+        raise NotImplementedError(f"{self.__class__.__name__} doesn't implement rmatmat")
 
     def solve(self, b: torch.Tensor) -> torch.Tensor:
         raise NotImplementedError(f"{self.__class__.__name__} doesn't implement solve")
 
+    def solve_plus_diag(self, b: torch.Tensor, diag: int | float | torch.Tensor) -> torch.Tensor:
+        return self.add_diagonal(diag).solve(b)
+
     def solve_bounded(self, b: torch.Tensor, bound:float, ord:float=2) -> torch.Tensor:
         """solve with a norm bound on x"""
         raise NotImplementedError(f"{self.__class__.__name__} doesn't implement solve_bounded")
 
-    def update(self, *args, **kwargs) -> None:
-        raise NotImplementedError(f"{self.__class__.__name__} doesn't implement update")
+    # def update(self, *args, **kwargs) -> None:
+    #     raise NotImplementedError(f"{self.__class__.__name__} doesn't implement update")
 
     def add(self, x: torch.Tensor) -> "LinearOperator":
         raise NotImplementedError(f"{self.__class__.__name__} doesn't implement add")
@@ -129,8 +137,8 @@ class Dense(LinearOperator):
     def matvec(self, x): return self.A.mv(x)
     def rmatvec(self, x): return self.A.mH.mv(x)
 
-    def matmat(self, x): return Dense(self.A.mm(x))
-    def rmatmat(self, x): return Dense(self.A.mH.mm(x))
+    def matmat(self, X): return Dense(self.A.mm(X))
+    def rmatmat(self, X): return Dense(self.A.mH.mm(X))
 
     def solve(self, b): return _solve(self.A, b)
 
@@ -146,6 +154,12 @@ class Dense(LinearOperator):
     def is_dense(self): return True
     def transpose(self): return Dense(self.A.mH)
 
+class SPD(Dense):
+    def solve(self, b: torch.Tensor):
+        L, info = torch.linalg.cholesky_ex(self.A) # pylint:disable=not-callable
+        return torch.cholesky_solve(b.unsqueeze(-1), L).squeeze(-1)
+
+
 class DenseInverse(LinearOperator):
     """Represents inverse of a dense matrix A."""
     def __init__(self, A_inv: torch.Tensor):
@@ -156,8 +170,8 @@ class DenseInverse(LinearOperator):
     def matvec(self, x): return _solve(self.A_inv, x) # pylint:disable=not-callable
     def rmatvec(self, x): return _solve(self.A_inv.mH, x) # pylint:disable=not-callable
 
-    def matmat(self, x): return Dense(_solve(self.A_inv, x)) # pylint:disable=not-callable
-    def rmatmat(self, x): return Dense(_solve(self.A_inv.mH, x)) # pylint:disable=not-callable
+    def matmat(self, X): return Dense(_solve(self.A_inv, X)) # pylint:disable=not-callable
+    def rmatmat(self, X): return Dense(_solve(self.A_inv.mH, X)) # pylint:disable=not-callable
 
     def solve(self, b): return self.A_inv.mv(b)
 
@@ -190,8 +204,8 @@ class Diagonal(LinearOperator):
     def matvec(self, x): return self.A * x
     def rmatvec(self, x): return self.A * x
 
-    def matmat(self, x): return Dense(x * self.A.unsqueeze(-1))
-    def rmatmat(self, x): return Dense(x * self.A.unsqueeze(-1))
+    def matmat(self, X): return Dense(X * self.A.unsqueeze(-1))
+    def rmatmat(self, X): return Dense(X * self.A.unsqueeze(-1))
 
     def solve(self, b): return b/self.A
 
@@ -221,8 +235,8 @@ class ScaledIdentity(LinearOperator):
     def matvec(self, x): return x * self.s
     def rmatvec(self, x): return x * self.s
 
-    def matmat(self, x): return Dense(x * self.s)
-    def rmatmat(self, x): return Dense(x * self.s)
+    def matmat(self, X): return Dense(X * self.s)
+    def rmatmat(self, X): return Dense(X * self.s)
 
     def solve(self, b): return b / self.s
     def solve_bounded(self, b, bound, ord = 2):
@@ -263,6 +277,7 @@ class ScaledIdentity(LinearOperator):
     def is_dense(self): return False
     def transpose(self): return ScaledIdentity(self.s, shape=self.shape, device=self.device, dtype=self.dtype)
 
+
 class AtA(LinearOperator):
     def __init__(self, A: torch.Tensor):
         self.A = A
@@ -270,8 +285,8 @@ class AtA(LinearOperator):
     def matvec(self, x): return self.A.mH.mv(self.A.mv(x))
     def rmatvec(self, x): return self.matvec(x)
 
-    def matmat(self, x): return Dense(torch.linalg.multi_dot([self.A.mH, self.A, x])) # pylint:disable=not-callable
-    def rmatmat(self, x): return Dense(torch.linalg.multi_dot([self.A.mH, self.A, x])) # pylint:disable=not-callable
+    def matmat(self, X): return Dense(torch.linalg.multi_dot([self.A.mH, self.A, X])) # pylint:disable=not-callable
+    def rmatmat(self, X): return Dense(torch.linalg.multi_dot([self.A.mH, self.A, X])) # pylint:disable=not-callable
 
     def is_dense(self): return False
     def to_tensor(self): return self.A.mH @ self.A
@@ -283,7 +298,27 @@ class AtA(LinearOperator):
         return Dense(self.to_tensor() + torch.diag_embed(x))
 
     def solve(self, b):
-        return Dense(self.to_tensor()).solve(b)
+        *_, n, m = self.A.shape
+        if n >= m: return Dense(self.to_tensor()).solve(b)
+
+        A = self.A
+        C = A @ A.mH # (n, n), SPD
+        L, info = torch.linalg.cholesky_ex(C) # pylint:disable=not-callable
+        z = torch.cholesky_solve((A @ b).unsqueeze(-1), L).squeeze(-1)
+        return A.mH @ z
+
+    def solve_plus_diag(self, b, diag):
+        *_, n, m = self.A.shape
+        if (n >= m) or (isinstance(diag, torch.Tensor) and diag.numel() > 1):
+            return Dense(self.to_tensor()).solve_plus_diag(b, diag)
+
+        A = self.A
+        I = torch.eye(A.size(-2), device=A.device, dtype=A.dtype)
+
+        C = (A @ A.mH).add_(I.mul_(diag)) # (n, n), SPD
+        L, info = torch.linalg.cholesky_ex(C + I.mul_(diag)) # pylint:disable=not-callable
+        z = torch.cholesky_solve((A @ b).unsqueeze(-1), L).squeeze(-1)
+        return (1 / diag) * (b - A.mH @ z)
 
     def inv(self):
         return Dense(self.to_tensor()).inv()
@@ -295,35 +330,98 @@ class AtA(LinearOperator):
         n = self.A.size(1)
         return (n,n)
 
-class AAT(LinearOperator):
+class AAt(AtA):
     def __init__(self, A: torch.Tensor):
-        self.A = A
+        super().__init__(A.mH)
 
-    def matvec(self, x): return self.A.mv(self.A.mH.mv(x))
-    def rmatvec(self, x): return self.matvec(x)
+class Sketched(LinearOperator):
+    """A projected by sketching matrix S, representing the operator S @ A_proj @ S.T.
+
+    Where A is (n, n) and S is (n, sketch_size).
+    """
+    def __init__(self, S: torch.Tensor, A_proj: torch.Tensor):
+        self.S = S
+        self.A_proj = A_proj
+        self.device = self.A_proj.device; self.dtype = self.A_proj.dtype
+
+    def matvec(self, x):
+        x_proj = self.S.T @ x
+        Ax_proj = self.A_proj @ x_proj
+        return self.S @ Ax_proj
+
+    def rmatvec(self, x):
+        x_proj = self.S.T @ x
+        ATx_proj = self.A_proj.mH @ x_proj
+        return self.S @ ATx_proj
+
+
+    def matmat(self, X): return Dense(torch.linalg.multi_dot([self.S, self.A_proj, self.S.T, X])) # pylint:disable=not-callable
+    def rmatmat(self, X): return Dense(torch.linalg.multi_dot([self.S, self.A_proj.mH, self.S.T, X])) # pylint:disable=not-callable
 
-    def matmat(self, x): return Dense(torch.linalg.multi_dot([self.A, self.A.mH, x])) # pylint:disable=not-callable
-    def rmatmat(self, x): return Dense(torch.linalg.multi_dot([self.A, self.A.mH, x])) # pylint:disable=not-callable
 
     def is_dense(self): return False
-    def to_tensor(self): return self.A @ self.A.mH
-    def transpose(self): return AAT(self.A)
+    def to_tensor(self): return self.S @ self.A_proj @ self.S.T
+    def transpose(self): return Sketched(self.S, self.A_proj.mH)
 
     def add_diagonal(self, x):
+        """this doesn't correspond to adding diagonal to A, however it still works for LM etc."""
         if isinstance(x, torch.Tensor) and x.numel() <= 1: x = x.item()
-        if isinstance(x, (int,float)): x = torch.full((self.shape[0],), fill_value=x, device=self.A.device, dtype=self.A.dtype)
-        return Dense(self.to_tensor() + torch.diag_embed(x))
+        if isinstance(x, (int,float)): x = torch.full((self.A_proj.shape[0],), fill_value=x, device=self.A_proj.device, dtype=self.A_proj.dtype)
+        return Sketched(S=self.S, A_proj=self.A_proj + x.diag_embed())
 
     def solve(self, b):
-        return Dense(self.to_tensor()).solve(b)
+        return self.S @ torch.linalg.lstsq(self.A_proj, self.S.T @ b).solution # pylint:disable=not-callable
 
     def inv(self):
-        return Dense(self.to_tensor()).inv()
-
-    def diagonal(self):
-        return self.A.pow(2).sum(0)
+        return Sketched(S=self.S, A_proj=torch.linalg.pinv(self.A_proj)) # pylint:disable=not-callable
 
     def size(self):
-        n = self.A.size(1)
+        n = self.S.size(0)
         return (n,n)
 
+
+class Eigendecomposition(LinearOperator):
+    """A represented as Q L Q^H. If A is (n,n), then Q is (n, rank); L is a vector - diagonal of (rank, rank)"""
+    def __init__(self, L: torch.Tensor, Q: torch.Tensor, use_nystrom: bool = True):
+        self.L = L
+        self.Q = Q
+        self.use_nystrom = use_nystrom
+        self.device = self.L.device; self.dtype = self.L.dtype
+
+    def matvec(self, x):
+        return self.Q @ ((self.Q.mH @ x) * self.L)
+
+    def rmatvec(self, x):
+        return self.matvec(x)
+
+    def matmat(self, X):
+        return Dense(self.Q @ (self.L[:, None] * (self.Q.mH @ X)))
+
+    def rmatmat(self, X):
+        return self.matmat(X)
+
+    def is_dense(self): return False
+    def to_tensor(self): return self.Q @ self.L.diag_embed() @ self.Q.mH
+    def transpose(self): return Eigendecomposition(L=self.L, Q=self.Q)
+
+    def add_diagonal(self, x):
+        """this doesn't correspond to adding diagonal to A, however it still works for LM etc."""
+        if isinstance(x, torch.Tensor) and x.numel() > 1:
+            raise RuntimeError("Eigendecomposition linear operator doesn't support add_diagonal with a vector diag")
+
+        return Eigendecomposition(L=self.L + x, Q = self.Q)
+
+    def solve(self, b):
+        return self.Q @ ((self.Q.mH @ b) / self.L)
+
+    def solve_plus_diag(self, b, diag):
+        if isinstance(diag, torch.Tensor) and diag.numel() > 1: return super().solve_plus_diag(b, diag)
+        if not self.use_nystrom: return super().solve_plus_diag(b, diag)
+        return nystrom_sketch_and_solve(L=self.L, Q=self.Q, b=b, reg=float(diag))
+
+    def inv(self):
+        return Eigendecomposition(L=1 / self.L, Q = self.Q)
+
+    def size(self):
+        n = self.Q.size(0)
+        return (n,n)

torchzero/linalg/matrix_power.py

@@ -0,0 +1,28 @@
+from typing import Literal
+import warnings
+from collections.abc import Callable
+
+import torch
+from . import torch_linalg
+def matrix_power_eigh(A: torch.Tensor, power:float, abs:bool=False):
+    """this is faster than SVD but only for positive semi-definite symmetric matrices
+    (covariance matrices are always SPD)"""
+
+    L, Q = torch_linalg.eigh(A, retry_float64=True) # pylint:disable=not-callable
+    if abs: L.abs_()
+    if power % 2 != 0: L.clip_(min = torch.finfo(A.dtype).tiny * 2)
+    return (Q * L.pow_(power).unsqueeze(-2)) @ Q.mH
+
+
+def matrix_power_svd(A: torch.Tensor, power: float) -> torch.Tensor:
+    """for any symmetric matrix"""
+    U, S, Vh = torch_linalg.svd(A, full_matrices=False, retry_float64=True) # pylint:disable=not-callable
+    if power % 2 != 0: S.clip_(min = torch.finfo(A.dtype).tiny * 2)
+    return (U * S.pow_(power).unsqueeze(-2)) @ Vh
+
+MatrixPowerMethod = Literal["eigh", "eigh_abs", "svd"]
+def matrix_power(A: torch.Tensor, power: float, method: MatrixPowerMethod = "eigh_abs") -> torch.Tensor:
+    if method == "eigh": return matrix_power_eigh(A, power)
+    if method == "eigh_abs": return matrix_power_eigh(A, power, abs=True)
+    if method == "svd": return matrix_power_svd(A, power)
+    raise ValueError(method)

torchzero/linalg/orthogonalize.py

@@ -0,0 +1,95 @@
+from typing import Literal
+import torch
+
+from ..utils.compile import allow_compile
+from . import torch_linalg
+
+# zeropower_via_newtonschulz5 from:
+# https://github.com/KellerJordan/Muon/blob/master/muon.py
+# and
+# https://github.com/HomebrewML/HeavyBall/blob/main/heavyball/utils.py#L452
+_NS_COEFFS = (
+    (4.0848, -6.8946, 2.9270),
+    (3.9505, -6.3029, 2.6377),
+    (3.7418, -5.5913, 2.3037),
+    (2.8769, -3.1427, 1.2046),
+    (2.8366, -3.0525, 1.2012)
+)
+
+@allow_compile
+def zeropower_via_newtonschulz5(G: torch.Tensor, coeffs=_NS_COEFFS) -> torch.Tensor:
+    """
+    Applies to last 2 dims - so usually reverse_dims should be applied to G before and after.
+
+    Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
+    quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose
+    of minimizing steps, it turns out to be empirically effective to keep increasing the slope at
+    zero even beyond the point where the iteration no longer converges all the way to one everywhere
+    on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T
+    where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model
+    performance at all relative to UV^T, where USV^T = G is the SVD.
+    """
+    assert G.ndim >= 2 # batched Muon implementation by @scottjmaddox, and put into practice in the record by @YouJiacheng
+
+    X = G.bfloat16()
+    if G.size(-2) > G.size(-1):
+        X = X.mT
+
+    # Ensure spectral norm is at most 1
+    X = X / (X.norm(dim=(-2, -1), keepdim=True).clip(min=torch.finfo(X.dtype).tiny * 2))
+
+    # Perform the NS iterations
+    for a,b,c in coeffs:
+        A = X @ X.mT
+        B = b * A + c * A @ A # quintic computation strategy adapted from suggestion by @jxbz, @leloykun, and @YouJiacheng
+        X = a * X + B @ X
+
+    if G.size(-2) > G.size(-1):
+        X = X.mT
+
+    return X.to(G.dtype)
+
+# code from https://github.com/MarkTuddenham/Orthogonal-Optimisers.
+# Tuddenham, M., Prügel-Bennett, A., & Hare, J. (2022).
+# Orthogonalising gradients to speed up neural network optimisation. arXiv preprint arXiv:2202.07052.
+def zeropower_via_svd(A: torch.Tensor) -> torch.Tensor:
+    """
+    Applies to first 2 dims and isn't batched - rest of dimensions are flattened.
+    """
+    try:
+        U, S, Vt = torch_linalg.svd(A, full_matrices=False, retry_float64=True) # pylint:disable=not-callable
+    except torch.linalg.LinAlgError:
+        U, S, Vt = torch.svd_lowrank(A, q=1, M=1e-4 * A.mean() * torch.rand_like(A))
+
+    return  U @ Vt
+
+def zeropower_via_eigh(A: torch.Tensor) -> torch.Tensor:
+    """
+    Only SPD and I need to check if I apply those to SPD because this is better than SVD.
+    """
+    L, Q = torch_linalg.eigh(A, retry_float64=True)
+    return  Q @ Q.mH
+
+
+def orthogonalize_via_qr(A: torch.Tensor):
+    *_, m, n = A.shape
+    T = False
+    if m < n:
+        T = True
+        m,n = n,m
+        A = A.mH
+
+    Q = torch_linalg.qr(A, mode='reduced', retry_float64=True).Q
+
+    if T:
+        Q = Q.mH
+
+    return Q
+
+OrthogonalizeMethod = Literal["newtonschulz", "svd", "qr"]
+def orthogonalize(A: torch.Tensor, method: OrthogonalizeMethod = "newtonschulz") -> torch.Tensor:
+    if method == "newtonschulz": return zeropower_via_newtonschulz5(A)
+    if method == "svd": return zeropower_via_svd(A)
+    if method == "qr": return orthogonalize_via_qr(A)
+    if method == "eigh": return zeropower_via_eigh(A)
+    raise ValueError(method)

torchzero/{utils/linalg → linalg}/qr.py

@@ -1,8 +1,9 @@
 from typing import Literal
 import torch
-from ..compile import enable_compilation
+from ..utils.compile import allow_compile
 
 # reference - https://www.cs.cornell.edu/~bindel/class/cs6210-f09/lec18.pdf
+@allow_compile
 def _get_w_tau(R: torch.Tensor, i: int, eps: float):
     R_ii = R[...,i,i]
     R_below = R[...,i:,i]
@@ -17,6 +18,7 @@ def _get_w_tau(R: torch.Tensor, i: int, eps: float):
     tau = torch.where(degenerate, 1, tau)
     return w, tau
 
+@allow_compile
 def _qr_householder_complete(A:torch.Tensor):
     *b,m,n = A.shape
     k = min(m,n)
@@ -33,6 +35,7 @@ def _qr_householder_complete(A:torch.Tensor):
 
     return Q, R
 
+@allow_compile
 def _qr_householder_reduced(A:torch.Tensor):
     *b,m,n = A.shape
     k = min(m,n)
@@ -64,7 +67,6 @@ def _qr_householder_reduced(A:torch.Tensor):
 
     return Q, R
 
-# @enable_compilation
 def qr_householder(A:torch.Tensor, mode: Literal['complete', 'reduced'] = 'reduced'):
     """an attempt at making QR decomposition for very tall and thin matrices that doesn't freeze, but it is around n_cols times slower than torch.linalg.qr, but compilation makes it faster, but it has to recompile when processing different shapes"""
     if mode == 'reduced': return _qr_householder_reduced(A)