torchzero 0.1.8__py3-none-any.whl → 0.3.1__py3-none-any.whl

torchzero/utils/linalg/matrix_funcs.py

@@ -0,0 +1,87 @@
+import warnings
+from collections.abc import Callable
+
+import torch
+
+def eigvals_func(A: torch.Tensor, fn: Callable[[torch.Tensor], torch.Tensor]) -> torch.Tensor:
+    L, Q = torch.linalg.eigh(A) # pylint:disable=not-callable
+    L = fn(L)
+    return  (Q * L.unsqueeze(-2)) @ Q.mH
+
+def singular_vals_func(A: torch.Tensor, fn: Callable[[torch.Tensor], torch.Tensor]) -> torch.Tensor:
+    U, S, V = torch.linalg.svd(A) # pylint:disable=not-callable
+    S = fn(S)
+    return (U * S.unsqueeze(-2)) @ V.mT
+
+def matrix_power_eigh(A: torch.Tensor, pow:float):
+    L, Q = torch.linalg.eigh(A) # pylint:disable=not-callable
+    if pow % 2 != 0: L.clip_(min = torch.finfo(A.dtype).eps)
+    return (Q * L.pow(pow).unsqueeze(-2)) @ Q.mH
+
+
+def inv_sqrt_2x2(A: torch.Tensor, force_pd: bool=False) -> torch.Tensor:
+    """Inverse square root of a possibly batched 2x2 matrix using a general formula for 2x2 matrices so that this is way faster than torch linalg. I tried doing a hierarchical 2x2 preconditioning but it didn't work well."""
+    eps = torch.finfo(A.dtype).eps
+
+    a = A[..., 0, 0]
+    b = A[..., 0, 1]
+    c = A[..., 1, 0]
+    d = A[..., 1, 1]
+
+    det = (a * d).sub_(b * c)
+    trace = a + d
+
+    if force_pd:
+        # add smallest eigenvalue magnitude to diagonal to force PD
+        # could also abs or clip eigenvalues bc there is a formula for eigenvectors
+        term1 = trace/2
+        term2 = (trace.pow(2).div_(4).sub_(det)).clamp_(min=eps).sqrt_()
+        y1 = term1 + term2
+        y2 = term1 - term2
+        smallest_eigval = torch.minimum(y1, y2).neg_().clamp_(min=0) + eps
+        a = a+smallest_eigval
+        d = d+smallest_eigval
+
+        # recalculate det and trace witg new a and b
+        det = (a * d).sub_(b * c)
+        trace = a + d
+
+    s = (det.clamp(min=eps)).sqrt_()
+
+    tau_squared = trace + 2 * s
+    tau = (tau_squared.clamp(min=eps)).sqrt_()
+
+    denom = s * tau
+
+    coeff = (denom.clamp(min=eps)).reciprocal_().unsqueeze(-1).unsqueeze(-1)
+
+    row1 = torch.stack([d + s, -b], dim=-1)
+    row2 = torch.stack([-c, a + s], dim=-1)
+    M = torch.stack([row1, row2], dim=-2)
+
+    return coeff * M
+
+
+def x_inv(diag: torch.Tensor,antidiag: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """invert a matrix with diagonal and anti-diagonal non zero elements, with no checks that it is invertible"""
+    n = diag.shape[0]
+    if diag.dim() != 1 or antidiag.dim() != 1 or antidiag.shape[0] != n:
+        raise ValueError("Input tensors must be 1D and have the same size.")
+    if n == 0:
+        return torch.empty_like(diag), torch.empty_like(antidiag)
+
+    # opposite indexes
+    diag_rev = torch.flip(diag, dims=[0])
+    antidiag_rev = torch.flip(antidiag, dims=[0])
+
+    # determinants
+    # det_i = d[i] * d[n-1-i] - a[i] * a[n-1-i]
+    determinant_vec = diag * diag_rev - antidiag * antidiag_rev
+
+    # inverse diagonal elements: y_d[i] = d[n-1-i] / det_i
+    inv_diag_vec = diag_rev / determinant_vec
+
+    # inverse anti-diagonal elements: y_a[i] = -a[i] / det_i
+    inv_anti_diag_vec = -antidiag / determinant_vec
+
+    return inv_diag_vec, inv_anti_diag_vec

torchzero/utils/linalg/orthogonalize.py

@@ -0,0 +1,11 @@
+from typing import overload
+import torch
+from ..tensorlist import TensorList
+
+@overload
+def gram_schmidt(x: torch.Tensor, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: ...
+@overload
+def gram_schmidt(x: TensorList, y: TensorList) -> tuple[TensorList, TensorList]: ...
+def gram_schmidt(x, y):
+    """makes two orthogonal vectors, only y is changed"""
+    return x, y - (x*y) / ((x*x) + 1e-8)

torchzero/utils/linalg/qr.py

@@ -0,0 +1,71 @@
+from typing import Literal
+import torch
+from ..compile import enable_compilation
+
+# reference - https://www.cs.cornell.edu/~bindel/class/cs6210-f09/lec18.pdf
+def _get_w_tau(R: torch.Tensor, i: int, eps: float):
+    R_ii = R[...,i,i]
+    R_below = R[...,i:,i]
+    norm_x = torch.linalg.vector_norm(R_below, dim=-1) # pylint:disable=not-callable
+    degenerate = norm_x < eps
+    s = -torch.sign(R_ii)
+    u1 = R_ii - s*norm_x
+    u1 = torch.where(degenerate, 1, u1)
+    w = R_below / u1.unsqueeze(-1)
+    w[...,0] = 1
+    tau = -s*u1/norm_x
+    tau = torch.where(degenerate, 1, tau)
+    return w, tau
+
+def _qr_householder_complete(A:torch.Tensor):
+    *b,m,n = A.shape
+    k = min(m,n)
+    eps = torch.finfo(A.dtype).eps
+
+    Q = torch.eye(m, dtype=A.dtype, device=A.device).expand(*b, m, m).clone() # clone because expanded dims refer to same memory
+    R = A.clone()
+
+    for i in range(k):
+        w, tau = _get_w_tau(R, i, eps)
+
+        R[..., i:,:] -= (tau*w).unsqueeze(-1) @ (w.unsqueeze(-2) @ R[..., i:,:])
+        Q[..., :,i:] -= (Q[..., :,i:]@w).unsqueeze(-1) @ (tau*w).unsqueeze(-2)
+
+    return Q, R
+
+def _qr_householder_reduced(A:torch.Tensor):
+    *b,m,n = A.shape
+    k = min(m,n)
+    eps = torch.finfo(A.dtype).eps
+
+    R = A.clone()
+
+    ws:list = [None for _ in range(k)]
+    taus:list = [None for _ in range(k)]
+
+    for i in range(k):
+        w, tau = _get_w_tau(R, i, eps)
+
+        ws[i] = w
+        taus[i] = tau
+
+        if m - i > 0 :
+            R[...,  i:,:] -= (tau*w).unsqueeze(-1) @ (w.unsqueeze(-2) @ R[..., i:,:])
+            # Q[..., :,i:] -= (Q[..., :,i:]@w).unsqueeze(-1) @ (tau*w).unsqueeze(-2)
+
+    R = R[..., :k, :]
+    Q = torch.eye(m, k, dtype=A.dtype, device=A.device).expand(*b, m, k).clone()
+    for i in range(k - 1, -1, -1):
+        if m - i > 0:
+            w = ws[i]
+            tau = taus[i].unsqueeze(-1).unsqueeze(-1)
+            Q_below = Q[..., i:, :]
+            Q[..., i:, :] -= torch.linalg.multi_dot([tau * w.unsqueeze(-1), w.unsqueeze(-2), Q_below]) # pylint:disable=not-callable
+
+    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)
+    return _qr_householder_complete(A)

torchzero/utils/linalg/solve.py

@@ -0,0 +1,168 @@
+from collections.abc import Callable
+from typing import overload
+import torch
+
+from .. import TensorList, generic_zeros_like, generic_vector_norm, generic_numel, generic_randn_like, generic_eq
+
+@overload
+def cg(
+    A_mm: Callable[[torch.Tensor], torch.Tensor],
+    b: torch.Tensor,
+    x0_: torch.Tensor | None,
+    tol: float | None,
+    maxiter: int | None,
+    reg: float = 0,
+) -> torch.Tensor: ...
+@overload
+def cg(
+    A_mm: Callable[[TensorList], TensorList],
+    b: TensorList,
+    x0_: TensorList | None,
+    tol: float | None,
+    maxiter: int | None,
+    reg: float | list[float] | tuple[float] = 0,
+) -> TensorList: ...
+
+def cg(
+    A_mm: Callable,
+    b: torch.Tensor | TensorList,
+    x0_: torch.Tensor | TensorList | None,
+    tol: float | None,
+    maxiter: int | None,
+    reg: float | list[float] | tuple[float] = 0,
+):
+    def A_mm_reg(x): # A_mm with regularization
+        Ax = A_mm(x)
+        if not generic_eq(reg, 0): Ax += x*reg
+        return Ax
+
+    if maxiter is None: maxiter = generic_numel(b)
+    if x0_ is None: x0_ = generic_zeros_like(b)
+
+    x = x0_
+    residual = b - A_mm_reg(x)
+    p = residual.clone() # search direction
+    r_norm = generic_vector_norm(residual)
+    init_norm = r_norm
+    if tol is not None and r_norm < tol: return x
+    k = 0
+
+    while True:
+        Ap = A_mm_reg(p)
+        step_size = (r_norm**2) / p.dot(Ap)
+        x += step_size * p # Update solution
+        residual -= step_size * Ap # Update residual
+        new_r_norm = generic_vector_norm(residual)
+
+        k += 1
+        if tol is not None and new_r_norm <= tol * init_norm: return x
+        if k >= maxiter: return x
+
+        beta = (new_r_norm**2) / (r_norm**2)
+        p = residual + beta*p
+        r_norm = new_r_norm
+
+
+# https://arxiv.org/pdf/2110.02820 algorithm 2.1 apparently supposed to be diabolical
+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
+
+# this one works worse
+def nystrom_sketch_and_solve(
+    A_mm: Callable[[torch.Tensor], torch.Tensor],
+    b: torch.Tensor,
+    rank: int,
+    reg: float,
+    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,
+    )
+    b = b.unsqueeze(-1)
+    lambd += 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)
+    return (term1 + term2).squeeze(-1)
+
+# this one is insane
+def nystrom_pcg(
+    A_mm: Callable[[torch.Tensor], torch.Tensor],
+    b: torch.Tensor,
+    sketch_size: int,
+    reg: float,
+    x0_: torch.Tensor | None,
+    tol: float | None,
+    maxiter: int | 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
+
+    def A_mm_reg(x): # A_mm with regularization
+        Ax = A_mm(x)
+        if reg != 0: Ax += x*reg
+        return Ax
+
+    if maxiter is None: maxiter = b.numel()
+    if x0_ is None: x0_ = torch.zeros_like(b)
+
+    x = x0_
+    residual = b - A_mm_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
+    P_inv = term1 + term2
+    z = P_inv @ residual
+    p = z.clone() # search direction
+
+    init_norm = torch.linalg.vector_norm(residual) # pylint:disable=not-callable
+    if tol is not None and init_norm < tol: return x
+    k = 0
+    while True:
+        Ap = A_mm_reg(p)
+        rz = residual.dot(z)
+        step_size = rz / p.dot(Ap)
+        x += step_size * p
+        residual -= step_size * Ap
+
+        k += 1
+        if tol is not None and torch.linalg.vector_norm(residual) <= tol * init_norm: return x # pylint:disable=not-callable
+        if k >= maxiter: return x
+
+        z = P_inv @ residual
+        beta = residual.dot(z) / rz
+        p = z + p*beta

torchzero/utils/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/utils/numberlist.py

@@ -0,0 +1,132 @@
+"""A lightweight data type for a list of numbers (or anything else) with arithmetic overloads (using basic for-loops).
+Subclasses list so works with torch._foreach_xxx operations."""
+import builtins
+from collections.abc import Callable, Sequence, Iterable, Generator, Iterator
+import math
+import operator
+from typing import Any, Literal, TypedDict
+from typing_extensions import Self, TypeAlias, Unpack
+
+import torch
+from .python_tools import zipmap
+
+def _alpha_add(x, other, alpha):
+    return x + other * alpha
+
+def as_numberlist(x):
+    if isinstance(x, NumberList): return x
+    return NumberList(x)
+
+
+def maybe_numberlist(x):
+    if isinstance(x, (list,tuple)): return as_numberlist(x)
+    return x
+
+def _clamp(x,min,max):
+    if min is not None and x < min: return min
+    if max is not None and x > max: return max
+    return x
+
+class NumberList(list[int | float | Any]):
+    """List of python numbers.
+    Note that this only supports basic arithmetic operations that are overloaded.
+
+    Can't use a numpy array because _foreach methods do not work with it."""
+    # remove torch.Tensor from return values
+    # this is no longer necessary
+    # def __getitem__(self, i) -> Any:
+    #     return super().__getitem__(i)
+
+    # def __iter__(self) -> Iterator[Any]:
+    #     return super().__iter__()
+
+    def __add__(self, other: Any) -> Self: return self.add(other) # type:ignore
+    def __radd__(self, other: Any) -> Self: return self.add(other)
+
+    def __sub__(self, other: Any) -> Self: return self.sub(other)
+    def __rsub__(self, other: Any) -> Self: return self.sub(other).neg()
+
+    def __mul__(self, other: Any) -> Self: return self.mul(other) # type:ignore
+    def __rmul__(self, other: Any) -> Self: return self.mul(other) # type:ignore
+
+    def __truediv__(self, other: Any) -> Self: return self.div(other)
+    def __rtruediv__(self, other: Any):
+        if isinstance(other, (tuple,list)): return self.__class__(o / i for o, i in zip(self, other))
+        return self.__class__(other / i for i in self)
+
+    def __floordiv__(self, other: Any): return self.floor_divide(other)
+    def __mod__(self, other: Any): return self.remainder(other)
+
+
+    def __pow__(self, other: Any): return self.pow(other)
+    def __rpow__(self, other: Any): return self.rpow(other)
+
+    def __neg__(self): return self.neg()
+
+    def __eq__(self, other: Any): return self.eq(other) # type:ignore
+    def __ne__(self, other: Any): return self.ne(other) # type:ignore
+    def __lt__(self, other: Any): return self.lt(other) # type:ignore
+    def __le__(self, other: Any): return self.le(other) # type:ignore
+    def __gt__(self, other: Any): return self.gt(other) # type:ignore
+    def __ge__(self, other: Any): return self.ge(other) # type:ignore
+
+    def __invert__(self): return self.logical_not()
+
+    def __and__(self, other: Any): return self.logical_and(other)
+    def __or__(self, other: Any): return self.logical_or(other)
+    def __xor__(self, other: Any): return self.logical_xor(other)
+
+    def __bool__(self):
+        raise RuntimeError(f'Boolean value of {self.__class__.__name__} is ambiguous')
+
+    def zipmap(self, fn: Callable, other: Any | list | tuple, *args, **kwargs):
+        """If `other` is list/tuple, applies `fn` to this TensorList zipped with `other`.
+        Otherwise applies `fn` to this TensorList and `other`.
+        Returns a new TensorList with return values of the callable."""
+        return zipmap(self, fn, other, *args, **kwargs)
+
+    def zipmap_args(self, fn: Callable[..., Any], *others, **kwargs):
+        """If `args` is list/tuple, applies `fn` to this TensorList zipped with `others`.
+        Otherwise applies `fn` to this TensorList and `other`."""
+        others = [i if isinstance(i, (list, tuple)) else [i]*len(self) for i in others]
+        return self.__class__(fn(*z, **kwargs) for z in zip(self, *others))
+
+    # def _set_to_method_result_(self, method: str, *args, **kwargs):
+    #     """Sets each element of the tensorlist to the result of calling the specified method on the corresponding element.
+    #     This is used to support/mimic in-place operations, although I decided to remove them."""
+    #     res = getattr(self, method)(*args, **kwargs)
+    #     for i,v in enumerate(res): self[i] = v
+    #     return self
+
+    def add(self, other: Any, alpha: int | float = 1):
+        if alpha == 1: return self.zipmap(operator.add, other=other)
+        return self.zipmap(_alpha_add, other=other, alpha = alpha)
+
+    def sub(self, other: Any, alpha: int | float = 1):
+        if alpha == 1: return self.zipmap(operator.sub, other=other)
+        return self.zipmap(_alpha_add, other=other, alpha = -alpha)
+
+    def neg(self): return self.__class__(-i for i in self)
+    def mul(self, other: Any): return self.zipmap(operator.mul, other=other)
+    def div(self, other: Any) -> Self: return self.zipmap(operator.truediv, other=other)
+    def pow(self, exponent: Any): return self.zipmap(math.pow, other=exponent)
+    def floor_divide(self, other: Any): return self.zipmap(operator.floordiv, other=other)
+    def remainder(self, other: Any): return self.zipmap(operator.mod, other=other)
+    def rpow(self, other: Any): return self.zipmap(lambda x,y: y**x, other=other)
+
+    def fill_none(self, value):
+        if isinstance(value, (list,tuple)): return self.__class__(v if s is None else s for s, v in zip(self, value))
+        return self.__class__(value if s is None else s for s in self)
+
+    def logical_not(self): return self.__class__(not i for i in self)
+    def logical_and(self, other: Any): return self.zipmap(operator.and_, other=other)
+    def logical_or(self, other: Any): return self.zipmap(operator.or_, other=other)
+    def logical_xor(self, other: Any): return self.zipmap(operator.xor, other=other)
+
+    def map(self, fn: Callable[..., torch.Tensor], *args, **kwargs):
+        """Applies `fn` to all elements of this TensorList
+        and returns a new TensorList with return values of the callable."""
+        return self.__class__(fn(i, *args, **kwargs) for i in self)
+
+    def clamp(self, min=None, max=None):
+        return self.zipmap_args(_clamp, min, max)

torchzero/utils/ops.py

@@ -0,0 +1,10 @@
+import torch
+
+
+def maximum_(input:torch.Tensor, other: torch.Tensor):
+    """in-place maximum"""
+    return torch.maximum(input, other, out = input)
+
+def where_(input: torch.Tensor, condition: torch.Tensor, other: torch.Tensor):
+    """in-place where"""
+    return torch.where(condition, input, other, out = input)