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

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

@@ -1,3 +1,4 @@
+# pylint: disable = non-ascii-name
 # pyright: reportArgumentType=false
 import math
 from collections import deque
@@ -5,8 +6,8 @@ from collections.abc import Callable
 from typing import Any, NamedTuple, overload
 
 import torch
-
-from .. import (
+from .linalg_utils import mm
+from ..utils import (
     TensorList,
     generic_eq,
     generic_finfo_tiny,
@@ -15,88 +16,73 @@ from .. import (
     generic_zeros_like,
 )
 
-
-def _make_A_mm_reg(A_mm: Callable, reg):
-    def A_mm_reg(x): # A_mm with regularization
-        Ax = A_mm(x)
+def _make_A_mv_reg(A_mv: Callable, reg):
+    def A_mv_reg(x): # A_mm with regularization
+        Ax = A_mv(x)
         if not generic_eq(reg, 0): Ax += x*reg
         return Ax
-    return A_mm_reg
+    return A_mv_reg
 
 def _identity(x): return x
 
-
-# https://arxiv.org/pdf/2110.02820
-def nystrom_approximation(
-    A_mm: Callable[[torch.Tensor], torch.Tensor],
-    ndim: int,
-    rank: int,
-    device,
-    dtype = torch.float32,
-    generator = None,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    omega = torch.randn((ndim, rank), device=device, dtype=dtype, generator=generator) # Gaussian test matrix
-    omega, _ = torch.linalg.qr(omega) # Thin QR decomposition # pylint:disable=not-callable
-
-    # Y = AΩ
-    Y = torch.stack([A_mm(col) for col in omega.unbind(-1)], -1) # rank matvecs
-    v = torch.finfo(dtype).eps * torch.linalg.matrix_norm(Y, ord='fro') # Compute shift # pylint:disable=not-callable
-    Yv = Y + v*omega # Shift for stability
-    C = torch.linalg.cholesky_ex(omega.mT @ Yv)[0] # pylint:disable=not-callable
-    B = torch.linalg.solve_triangular(C, Yv.mT, upper=False, unitriangular=False).mT # pylint:disable=not-callable
-    U, S, _ = torch.linalg.svd(B, full_matrices=False) # pylint:disable=not-callable
-    lambd = (S.pow(2) - v).clip(min=0) #Remove shift, compute eigs
-    return U, lambd
-
+# TODO this is used in NystromSketchAndSolve
+# I need to add alternative to it where it just shifts eigenvalues by reg and uses their reciprocal
 def nystrom_sketch_and_solve(
-    A_mm: Callable[[torch.Tensor], torch.Tensor],
+    L: torch.Tensor,
+    Q: torch.Tensor,
     b: torch.Tensor,
-    rank: int,
     reg: float = 1e-3,
-    generator=None,
 ) -> torch.Tensor:
-    U, lambd = nystrom_approximation(
-        A_mm=A_mm,
-        ndim=b.size(-1),
-        rank=rank,
-        device=b.device,
-        dtype=b.dtype,
-        generator=generator,
-    )
+    """Solves (Q diag(L) Q.T + reg*I)x = b. Becomes super unstable with reg smaller than like 1e-5.
+
+    Args:
+        L (torch.Tensor): eigenvalues, like from ``nystrom_approximation``
+        Q (torch.Tensor): eigenvectors, like from ``nystrom_approximation``
+        b (torch.Tensor): right hand side
+        reg (float, optional): regularization. Defaults to 1e-3.
+    """
+
     b = b.unsqueeze(-1)
-    lambd += reg
+    L += reg
     # x = (A + μI)⁻¹ b
-    # (A + μI)⁻¹ = U(Λ + μI)⁻¹Uᵀ + (1/μ)(b - UUᵀ)
-    # x = U(Λ + μI)⁻¹Uᵀb + (1/μ)(b - UUᵀb)
-    Uᵀb = U.T @ b
-    term1 = U @ ((1/lambd).unsqueeze(-1) * Uᵀb)
-    term2 = (1.0 / reg) * (b - U @ Uᵀb)
+    # (A + μI)⁻¹ = Q(L + μI)⁻¹Qᵀ + (1/μ)(b - QQᵀ)
+    # x = Q(L + μI)⁻¹Qᵀb + (1/μ)(b - QQᵀb)
+    Qᵀb = Q.T @ b
+    term1 = Q @ ((1/L).unsqueeze(-1) * Qᵀb)
+    term2 = (1.0 / reg) * (b - Q @ Qᵀb)
     return (term1 + term2).squeeze(-1)
 
 def nystrom_pcg(
-    A_mm: Callable[[torch.Tensor], torch.Tensor],
+    L: torch.Tensor,
+    Q: torch.Tensor,
+    A_mv: Callable[[torch.Tensor], torch.Tensor],
     b: torch.Tensor,
-    sketch_size: int,
     reg: float = 1e-6,
     x0_: torch.Tensor | None = None,
-    tol: float | None = 1e-4,
+    tol: float | None = 1e-8,
     maxiter: int | None = None,
-    generator=None,
 ) -> torch.Tensor:
-    U, lambd = nystrom_approximation(
-        A_mm=A_mm,
-        ndim=b.size(-1),
-        rank=sketch_size,
-        device=b.device,
-        dtype=b.dtype,
-        generator=generator,
-    )
-    lambd += reg
+    """conjugate gradient preconditioned by nystrom approximation.
+
+    The preconditioner can be computed by one matrix-matrix multiplication with A.
+    If matrix-matrix is efficient, then this is good (e.g. batched hessian-vector products in pytorch)
+
+    Args:
+        L (torch.Tensor): eigenvalues of approximation of A, like from ``nystrom_approximation``
+        Q (torch.Tensor): eigenvectors of approximation of A, like from ``nystrom_approximation``
+        A_mv (Callable[[torch.Tensor], torch.Tensor]): mat-vec func with hessian
+        b (torch.Tensor): right hand side
+        reg (float, optional): regularization. Defaults to 1e-6.
+        x0_ (torch.Tensor | None, optional): initial guess (modified in-place). Defaults to None.
+        tol (float | None, optional): tolerance for convergence. Defaults to 1e-4.
+        maxiter (int | None, optional): maximum number of iterations. Defaults to None.
+    """
+    L += reg
     eps = torch.finfo(b.dtype).tiny * 2
     if tol is None: tol = eps
 
-    def A_mm_reg(x): # A_mm with regularization
-        Ax = A_mm(x)
+    def A_mv_reg(x): # A_mm with regularization
+        Ax = A_mv(x)
         if reg != 0: Ax += x*reg
         return Ax
 
@@ -104,10 +90,10 @@ def nystrom_pcg(
     if x0_ is None: x0_ = torch.zeros_like(b)
 
     x = x0_
-    residual = b - A_mm_reg(x)
+    residual = b - A_mv_reg(x)
     # z0 = P⁻¹ r0
-    term1 = lambd[...,-1] * U * (1/lambd.unsqueeze(-2)) @ U.mT
-    term2 = torch.eye(U.size(-2), device=U.device,dtype=U.dtype) - U@U.mT
+    term1 = L[...,-1] * Q * (1/L.unsqueeze(-2)) @ Q.mT
+    term2 = torch.eye(Q.size(-2), device=Q.device,dtype=Q.dtype) - Q@Q.mT
     P_inv = term1 + term2
     z = P_inv @ residual
     p = z.clone() # search direction
@@ -116,7 +102,7 @@ def nystrom_pcg(
     if init_norm < tol: return x
     k = 0
     while True:
-        Ap = A_mm_reg(p)
+        Ap = A_mv_reg(p)
         rz = residual.dot(z)
         step_size = rz / p.dot(Ap)
         x += step_size * p
@@ -138,7 +124,7 @@ def _safe_clip(x: torch.Tensor):
     if x.abs() < eps: return x.new_full(x.size(), eps).copysign(x)
     return x
 
-def _trust_tau(x,d,trust_radius):
+def _trust_tau(x, d, trust_radius):
     xx = x.dot(x)
     xd = x.dot(d)
     dd = _safe_clip(d.dot(d))
@@ -150,10 +136,10 @@ def _trust_tau(x,d,trust_radius):
 
 
 class CG:
-    """Conjugate gradient method.
+    """Conjugate gradient method optionally with norm constraint.
 
     Args:
-        A_mm (Callable[[torch.Tensor], torch.Tensor] | torch.Tensor): Callable that returns matvec ``Ax``.
+        A_mv (Callable[[torch.Tensor], torch.Tensor] | torch.Tensor): Callable that returns matvec ``Ax``.
         b (torch.Tensor): right hand side
         x0 (torch.Tensor | None, optional): initial guess, defaults to zeros. Defaults to None.
         tol (float | None, optional): tolerance for convergence. Defaults to 1e-8.
@@ -174,10 +160,10 @@ class CG:
     """
     def __init__(
         self,
-        A_mm: Callable,
+        A_mv: Callable,
         b: torch.Tensor | TensorList,
         x0: torch.Tensor | TensorList | None = None,
-        tol: float | None = 1e-4,
+        tol: float | None = 1e-8,
         maxiter: int | None = None,
         reg: float = 0,
         trust_radius: float | None = None,
@@ -187,7 +173,7 @@ class CG:
         P_mm: Callable | None = None,
 ):
         # --------------------------------- set attrs -------------------------------- #
-        self.A_mm = _make_A_mm_reg(A_mm, reg)
+        self.A_mv = _make_A_mv_reg(A_mv, reg)
         self.b = b
         if tol is None: tol = generic_finfo_tiny(b) * 2
         self.tol = tol
@@ -214,7 +200,7 @@ class CG:
             self.r = b
         else:
             self.x = x0
-            self.r = b - A_mm(self.x)
+            self.r = b - A_mv(self.x)
 
         self.z = self.P_mm(self.r)
         self.d = self.z
@@ -229,7 +215,7 @@ class CG:
         if self.iter >= self.maxiter:
             return x, True
 
-        Ad = self.A_mm(d)
+        Ad = self.A_mv(d)
         dAd = d.dot(Ad)
 
         # check negative curvature
@@ -289,7 +275,8 @@ class CG:
         return sol
 
 def find_within_trust_radius(history, trust_radius: float):
-    """find first ``x`` in history that exceeds trust radius, if no such ``x`` exists, returns ``None``"""
+    """find first ``x`` in history that exceeds trust radius and returns solution within,
+    if no such ``x`` exists, returns ``None``"""
     for x, x_norm, d in reversed(tuple(history)):
         if x_norm <= trust_radius:
             return _trust_tau(x, d, trust_radius)
@@ -306,7 +293,7 @@ class _TensorListSolution(NamedTuple):
 
 @overload
 def cg(
-    A_mm: Callable[[torch.Tensor], torch.Tensor],
+    A_mv: Callable[[torch.Tensor], torch.Tensor],
     b: torch.Tensor,
     x0: torch.Tensor | None = None,
     tol: float | None = 1e-8,
@@ -320,7 +307,7 @@ def cg(
 ) -> _TensorSolution: ...
 @overload
 def cg(
-    A_mm: Callable[[TensorList], TensorList],
+    A_mv: Callable[[TensorList], TensorList],
     b: TensorList,
     x0: TensorList | None = None,
     tol: float | None = 1e-8,
@@ -333,7 +320,7 @@ def cg(
     P_mm: Callable[[TensorList], TensorList] | None = None
 ) -> _TensorListSolution: ...
 def cg(
-    A_mm: Callable,
+    A_mv: Callable,
     b: torch.Tensor | TensorList,
     x0: torch.Tensor | TensorList | None = None,
     tol: float | None = 1e-8,
@@ -346,7 +333,7 @@ def cg(
     P_mm: Callable | None = None
 ):
     solver = CG(
-        A_mm=A_mm,
+        A_mv=A_mv,
         b=b,
         x0=x0,
         tol=tol,
@@ -370,10 +357,10 @@ def cg(
 # Liu, Yang, and Fred Roosta. "MINRES: From negative curvature detection to monotonicity properties." SIAM Journal on Optimization 32.4 (2022): 2636-2661.
 @overload
 def minres(
-    A_mm: Callable[[torch.Tensor], torch.Tensor] | torch.Tensor,
+    A_mv: Callable[[torch.Tensor], torch.Tensor] | torch.Tensor,
     b: torch.Tensor,
     x0: torch.Tensor | None = None,
-    tol: float | None = 1e-4,
+    tol: float | None = 1e-8,
     maxiter: int | None = None,
     reg: float = 0,
     npc_terminate: bool=True,
@@ -381,26 +368,27 @@ def minres(
 ) -> torch.Tensor: ...
 @overload
 def minres(
-    A_mm: Callable[[TensorList], TensorList],
+    A_mv: Callable[[TensorList], TensorList],
     b: TensorList,
     x0: TensorList | None = None,
-    tol: float | None = 1e-4,
+    tol: float | None = 1e-8,
     maxiter: int | None = None,
     reg: float | list[float] | tuple[float] = 0,
     npc_terminate: bool=True,
     trust_radius: float | None = None,
 ) -> TensorList: ...
 def minres(
-    A_mm,
+    A_mv,
     b,
     x0: torch.Tensor | TensorList | None = None,
-    tol: float | None = 1e-4,
+    tol: float | None = 1e-8,
     maxiter: int | None = None,
     reg: float | list[float] | tuple[float] = 0,
     npc_terminate: bool=True,
     trust_radius: float | None = None, #trust region is experimental
 ):
-    A_mm_reg = _make_A_mm_reg(A_mm, reg)
+    """MINRES (experimental)"""
+    A_mv_reg = _make_A_mv_reg(A_mv, reg)
     eps = math.sqrt(generic_finfo_tiny(b) * 2)
     if tol is None: tol = eps
 
@@ -409,7 +397,7 @@ def minres(
         R = b
         x0 = generic_zeros_like(b)
     else:
-        R = b - A_mm_reg(x0)
+        R = b - A_mv_reg(x0)
 
     X: Any = x0
     beta = b_norm = generic_vector_norm(b)
@@ -429,7 +417,7 @@ def minres(
 
     for _ in range(maxiter):
 
-        P = A_mm_reg(V)
+        P = A_mv_reg(V)
         alpha = V.dot(P)
         P -= beta*V_prev
         P -= alpha*V

torchzero/linalg/svd.py

@@ -0,0 +1,20 @@
+# import torch
+
+# # projected svd
+# # adapted from https://github.com/smortezavi/Randomized_SVD_GPU
+# def randomized_svd(M: torch.Tensor, k: int, driver=None):
+#     *_, m, n = M.shape
+#     transpose = False
+#     if m < n:
+#         transpose = True
+#         M = M.mT
+#         m,n = n,m
+
+#     rand_matrix = torch.randn(size=(n, k), device=M.device, dtype=M.dtype)
+#     Q, _ = torch.linalg.qr(M @ rand_matrix, mode='reduced') # pylint:disable=not-callable
+#     smaller_matrix = Q.mT @ M
+#     U_hat, s, V = torch.linalg.svd(smaller_matrix, driver=driver, full_matrices=False) # pylint:disable=not-callable
+#     U = Q @ U_hat
+
+#     if transpose: return V.mT, s, U.mT
+#     return U, s, V

torchzero/linalg/torch_linalg.py

@@ -0,0 +1,168 @@
+"""torch linalg with correct typing and retries in float64"""
+from typing import NamedTuple
+
+import torch
+
+
+def cholesky(A: torch.Tensor, *, upper=False, retry_float64:bool=False) -> torch.Tensor:
+    """A - SPD, returns lower triangular L such that ``A = L @ L.mH`` also can pass L to ``torch.cholesky_solve``"""
+    try:
+        return torch.linalg.cholesky(A, upper=upper) # pylint:disable=not-callable
+
+    except torch.linalg.LinAlgError as e:
+        if not retry_float64: raise e
+        dtype = A.dtype
+        if dtype == torch.float64: raise e
+        return cholesky(A.to(torch.float64), upper=upper, retry_float64=False).to(dtype)
+
+
+class _QRTuple(NamedTuple):
+    Q: torch.Tensor
+    R: torch.Tensor
+
+def qr(A: torch.Tensor, mode='reduced', retry_float64:bool=False) -> _QRTuple:
+    """A - any matrix ``(*, m, n)`` (for some reason sometimes it takes ages on some matrices)
+
+    ### Returns (if mode = "reduced"):
+
+    Q: ``(*, m, k)`` - orthogonal
+
+    R: ``(*, k, n)`` - upper triangular
+
+    where ``k = min(m,n)``
+    """
+    try:
+        return torch.linalg.qr(A, mode=mode) # pylint:disable=not-callable
+
+    except torch.linalg.LinAlgError as e:
+        if not retry_float64: raise e
+        dtype = A.dtype
+        if dtype == torch.float64: raise e
+        Q, R = qr(A.to(torch.float64), mode=mode, retry_float64=False)
+        return _QRTuple(Q=Q.to(dtype), R=R.to(dtype))
+
+def eigh(A: torch.Tensor, UPLO="L", retry_float64:bool=False) -> tuple[torch.Tensor, torch.Tensor]:
+    """A - symmetric, returns ``(L, Q)``, ``A = Q @ torch.diag(L) @ Q.mH``, this is faster than SVD"""
+    try:
+        return torch.linalg.eigh(A, UPLO=UPLO) # pylint:disable=not-callable
+
+    except torch.linalg.LinAlgError as e:
+        if not retry_float64: raise e
+        dtype = A.dtype
+        if dtype == torch.float64: raise e
+        L, Q = eigh(A.to(torch.float64), UPLO=UPLO, retry_float64=False)
+        return L.to(dtype), Q.to(dtype)
+
+
+
+class _SVDTuple(NamedTuple):
+    U: torch.Tensor
+    S: torch.Tensor
+    Vh: torch.Tensor
+
+def svd(A: torch.Tensor, full_matrices=True, driver=None, retry_float64:bool=False) -> _SVDTuple:
+    """A - any matrix ``(*, n, m)``, but slows down if A isn't well conditioned, ``A = U @ torch.diag(S) @ Vh``
+
+    Don't forget to set ``full_matrices=False``
+
+    ### Returns:
+
+    U: ``(*, m, m)`` or ``(*, m, k)`` - orthogonal
+
+    S: ``(*, k,)`` - singular values
+
+    V^H: ``(*, n, n)`` or ``(*, n, k)`` - orthogonal
+
+    where ``k = min(m,n)``
+
+    ### Drivers
+
+    drivers are only supported on CUDA so A is moved to CUDA by this function if needed
+
+    from docs:
+
+    If A is well-conditioned (its condition number is not too large), or you do not mind some precision loss.
+
+    For a general matrix: ‘gesvdj’ (Jacobi method)
+
+    If A is tall or wide (m >> n or m << n): ‘gesvda’ (Approximate method)
+
+    If A is not well-conditioned or precision is relevant: ‘gesvd’ (QR based)
+
+    By default (driver= None), we call ‘gesvdj’ and, if it fails, we fallback to ‘gesvd’.
+    """
+    # drivers are only for CUDA
+    # also the only one that doesn't freeze is ‘gesvda’
+    device=None
+    if driver is not None:
+        device = A.device
+        A = A.cuda()
+
+    try:
+        U, S, Vh = torch.linalg.svd(A, full_matrices=full_matrices, driver=driver) # pylint:disable=not-callable
+        if device is not None:
+            U = U.to(device); S = S.to(device); Vh = Vh.to(device)
+        return _SVDTuple(U=U, S=S, Vh=Vh)
+
+    except torch.linalg.LinAlgError as e:
+        if not retry_float64: raise e
+        dtype = A.dtype
+        if dtype == torch.float64: raise e
+        U, S, Vh = svd(A.to(torch.float64), full_matrices=full_matrices, driver=driver, retry_float64=False)
+        return _SVDTuple(U=U.to(dtype), S=S.to(dtype), Vh=Vh.to(dtype))
+
+def solve(A: torch.Tensor, B: torch.Tensor, left:bool=True, retry_float64:bool=False) -> torch.Tensor:
+    """I think this uses LU"""
+    try:
+        return torch.linalg.solve(A, B, left=left) # pylint:disable=not-callable
+
+    except torch.linalg.LinAlgError as e:
+        if not retry_float64: raise e
+        dtype = A.dtype
+        if dtype == torch.float64: raise e
+        return solve(A.to(torch.float64), B.to(torch.float64), left=left, retry_float64=False).to(dtype)
+
+class _SolveExTuple(NamedTuple):
+    result: torch.Tensor
+    info: int
+
+def solve_ex(A: torch.Tensor, B: torch.Tensor, left:bool=True, retry_float64:bool=False) -> _SolveExTuple:
+    """I think this uses LU"""
+    result, info = torch.linalg.solve_ex(A, B, left=left) # pylint:disable=not-callable
+
+    if info != 0:
+        if not retry_float64: return _SolveExTuple(result, info)
+        dtype = A.dtype
+        if dtype == torch.float64: return _SolveExTuple(result, info)
+        result, info = solve_ex(A.to(torch.float64), B.to(torch.float64), retry_float64=False)
+        return _SolveExTuple(result.to(dtype), info)
+
+    return _SolveExTuple(result, info)
+
+def inv(A: torch.Tensor, retry_float64:bool=False) -> torch.Tensor:
+    try:
+        return torch.linalg.inv(A) # pylint:disable=not-callable
+
+    except torch.linalg.LinAlgError as e:
+        if not retry_float64: raise e
+        dtype = A.dtype
+        if dtype == torch.float64: raise e
+        return inv(A.to(torch.float64), retry_float64=False).to(dtype)
+
+
+class _InvExTuple(NamedTuple):
+    inverse: torch.Tensor
+    info: int
+
+def inv_ex(A: torch.Tensor, *, check_errors=False, retry_float64:bool=False) -> _InvExTuple:
+    """this retries in float64 but on fail info will be not 0"""
+    inverse, info = torch.linalg.inv_ex(A, check_errors=check_errors) # pylint:disable=not-callable
+
+    if info != 0:
+        if not retry_float64: return _InvExTuple(inverse, info)
+        dtype = A.dtype
+        if dtype == torch.float64: return _InvExTuple(inverse, info)
+        inverse, info = inv_ex(A.to(torch.float64), retry_float64=False)
+        return _InvExTuple(inverse.to(dtype), info)
+
+    return _InvExTuple(inverse, info)

torchzero/modules/adaptive/__init__.py

@@ -12,7 +12,7 @@ from .lmadagrad import LMAdagrad
 from .lion import Lion
 from .mars import MARSCorrection
 from .matrix_momentum import MatrixMomentum
-from .msam import MSAM, MSAMObjective
+from .msam import MSAMMomentum, MSAM
 from .muon import DualNormCorrection, MuonAdjustLR, Orthogonalize, orthogonalize_grads_
 from .natural_gradient import NaturalGradient
 from .orthograd import OrthoGrad, orthograd_