torchzero 0.3.9__py3-none-any.whl → 0.3.11__py3-none-any.whl

torchzero/utils/linalg/solve.py

@@ -1,12 +1,41 @@
+# pyright: reportArgumentType=false
 from collections.abc import Callable
-from typing import overload
+from typing import Any, overload
+
 import torch
 
-from .. import TensorList, generic_zeros_like, generic_vector_norm, generic_numel, generic_randn_like, generic_eq
+from .. import (
+    TensorList,
+    generic_eq,
+    generic_finfo_eps,
+    generic_numel,
+    generic_randn_like,
+    generic_vector_norm,
+    generic_zeros_like,
+)
+
+
+def _make_A_mm_reg(A_mm: Callable | torch.Tensor, reg):
+    if callable(A_mm):
+        def A_mm_reg(x): # A_mm with regularization
+            Ax = A_mm(x)
+            if not generic_eq(reg, 0): Ax += x*reg
+            return Ax
+        return A_mm_reg
+
+    if not isinstance(A_mm, torch.Tensor): raise TypeError(type(A_mm))
+
+    def Ax_reg(x): # A_mm with regularization
+        if A_mm.ndim == 1: Ax = A_mm * x
+        else: Ax = A_mm @ x
+        if reg != 0: Ax += x*reg
+        return Ax
+    return Ax_reg
+
 
 @overload
 def cg(
-    A_mm: Callable[[torch.Tensor], torch.Tensor],
+    A_mm: Callable[[torch.Tensor], torch.Tensor] | torch.Tensor,
     b: torch.Tensor,
     x0_: torch.Tensor | None = None,
     tol: float | None = 1e-4,
@@ -24,17 +53,17 @@ def cg(
 ) -> TensorList: ...
 
 def cg(
-    A_mm: Callable,
+    A_mm: Callable | torch.Tensor,
     b: torch.Tensor | TensorList,
     x0_: torch.Tensor | TensorList | None = None,
     tol: float | None = 1e-4,
     maxiter: int | None = 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
+    A_mm_reg = _make_A_mm_reg(A_mm, reg)
+    eps = generic_finfo_eps(b)
+
+    if tol is None: tol = eps
 
     if maxiter is None: maxiter = generic_numel(b)
     if x0_ is None: x0_ = generic_zeros_like(b)
@@ -44,9 +73,10 @@ def cg(
     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
+    if r_norm < tol: return x
     k = 0
 
+
     while True:
         Ap = A_mm_reg(p)
         step_size = (r_norm**2) / p.dot(Ap)
@@ -55,7 +85,7 @@ def cg(
         new_r_norm = generic_vector_norm(residual)
 
         k += 1
-        if tol is not None and new_r_norm <= tol * init_norm: return x
+        if new_r_norm <= tol * init_norm: return x
         if k >= maxiter: return x
 
         beta = (new_r_norm**2) / (r_norm**2)
@@ -131,6 +161,8 @@ def nystrom_pcg(
         generator=generator,
     )
     lambd += reg
+    eps = torch.finfo(b.dtype).eps ** 2
+    if tol is None: tol = eps
 
     def A_mm_reg(x): # A_mm with regularization
         Ax = A_mm(x)
@@ -150,7 +182,7 @@ def nystrom_pcg(
     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
+    if init_norm < tol: return x
     k = 0
     while True:
         Ap = A_mm_reg(p)
@@ -160,10 +192,217 @@ def nystrom_pcg(
         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 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
 
+
+def _safe_clip(x: torch.Tensor):
+    """makes sure scalar tensor x is not smaller than epsilon"""
+    assert x.numel() == 1, x.shape
+    eps = torch.finfo(x.dtype).eps
+    if x.abs() < eps: return x.new_full(x.size(), eps).copysign(x)
+    return x
+
+def _trust_tau(x,d,trust_region):
+    xx = x.dot(x)
+    xd = x.dot(d)
+    dd = _safe_clip(d.dot(d))
+
+    rad = (xd**2 - dd * (xx - trust_region**2)).clip(min=0).sqrt()
+    tau = (-xd + rad) / dd
+
+    return x + tau * d
+
+
+@overload
+def steihaug_toint_cg(
+    A_mm: Callable[[torch.Tensor], torch.Tensor] | torch.Tensor,
+    b: torch.Tensor,
+    trust_region: float,
+    x0: torch.Tensor | None = None,
+    tol: float | None = 1e-4,
+    maxiter: int | None = None,
+    reg: float = 0,
+) -> torch.Tensor: ...
+@overload
+def steihaug_toint_cg(
+    A_mm: Callable[[TensorList], TensorList],
+    b: TensorList,
+    trust_region: float,
+    x0: TensorList | None = None,
+    tol: float | None = 1e-4,
+    maxiter: int | None = None,
+    reg: float | list[float] | tuple[float] = 0,
+) -> TensorList: ...
+def steihaug_toint_cg(
+    A_mm: Callable | torch.Tensor,
+    b: torch.Tensor | TensorList,
+    trust_region: float,
+    x0: torch.Tensor | TensorList | None = None,
+    tol: float | None = 1e-4,
+    maxiter: int | None = None,
+    reg: float | list[float] | tuple[float] = 0,
+):
+    """
+    Solution is bounded to have L2 norm no larger than :code:`trust_region`. If solution exceeds :code:`trust_region`, CG is terminated early, so it is also faster.
+    """
+    A_mm_reg = _make_A_mm_reg(A_mm, reg)
+
+    x = x0
+    if x is None: x = generic_zeros_like(b)
+    r = b
+    d = r.clone()
+
+    eps = generic_finfo_eps(b)**2
+    if tol is None: tol = eps
+
+    if generic_vector_norm(r) < tol:
+        return x
+
+    if maxiter is None:
+        maxiter = generic_numel(b)
+
+    for _ in range(maxiter):
+        Ad = A_mm_reg(d)
+
+        d_Ad = d.dot(Ad)
+        if d_Ad <= eps:
+            return _trust_tau(x, d, trust_region)
+
+        alpha = r.dot(r) / d_Ad
+        p_next = x + alpha * d
+
+        # check if the step exceeds the trust-region boundary
+        if generic_vector_norm(p_next) >= trust_region:
+            return _trust_tau(x, d, trust_region)
+
+        # update step, residual and direction
+        x = p_next
+        r_next = r - alpha * Ad
+
+        if generic_vector_norm(r_next) < tol:
+            return x
+
+        beta = r_next.dot(r_next) / r.dot(r)
+        d = r_next + beta * d
+        r = r_next
+
+    return x
+
+
+
+# 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,
+    b: torch.Tensor,
+    x0: torch.Tensor | None = None,
+    tol: float | None = 1e-4,
+    maxiter: int | None = None,
+    reg: float = 0,
+    npc_terminate: bool=True,
+    trust_region: float | None = None,
+) -> torch.Tensor: ...
+@overload
+def minres(
+    A_mm: Callable[[TensorList], TensorList],
+    b: TensorList,
+    x0: TensorList | None = None,
+    tol: float | None = 1e-4,
+    maxiter: int | None = None,
+    reg: float | list[float] | tuple[float] = 0,
+    npc_terminate: bool=True,
+    trust_region: float | None = None,
+) -> TensorList: ...
+def minres(
+    A_mm,
+    b,
+    x0: torch.Tensor | TensorList | None = None,
+    tol: float | None = 1e-4,
+    maxiter: int | None = None,
+    reg: float | list[float] | tuple[float] = 0,
+    npc_terminate: bool=True,
+    trust_region: float | None = None,
+):
+    A_mm_reg = _make_A_mm_reg(A_mm, reg)
+    eps = generic_finfo_eps(b)
+    if tol is None: tol = eps**2
+
+    if maxiter is None: maxiter = generic_numel(b)
+    if x0 is None:
+        R = b
+        x0 = generic_zeros_like(b)
+    else:
+        R = b - A_mm_reg(x0)
+
+    X: Any = x0
+    beta = b_norm = generic_vector_norm(b)
+    if b_norm < eps**2:
+        return generic_zeros_like(b)
+
+
+    V = b / beta
+    V_prev = generic_zeros_like(b)
+    D = generic_zeros_like(b)
+    D_prev = generic_zeros_like(b)
+
+    c = -1
+    phi = tau = beta
+    s = delta1 = e = 0
+
+
+    for _ in range(maxiter):
+
+        P = A_mm_reg(V)
+        alpha = V.dot(P)
+        P -= beta*V_prev
+        P -= alpha*V
+        beta = generic_vector_norm(P)
+
+        delta2 = c*delta1 + s*alpha
+        gamma1 = s*delta1 - c*alpha
+        e_next = s*beta
+        delta1 = -c*beta
+
+        cgamma1 = c*gamma1
+        if trust_region is not None and cgamma1 >= 0:
+            if npc_terminate: return _trust_tau(X, R, trust_region)
+            return _trust_tau(X, D, trust_region)
+
+        if npc_terminate and cgamma1 >= 0:
+            return R
+
+        gamma2 = (gamma1**2 + beta**2)**(1/2)
+
+        if abs(gamma2) <= eps: # singular system
+            # c=0; s=1; tau=0
+            if trust_region is None: return X
+            return _trust_tau(X, D, trust_region)
+
+        c = gamma1 / gamma2
+        s = beta/gamma2
+        tau = c*phi
+        phi = s*phi
+
+        D_prev = D
+        D = (V - delta2*D - e*D_prev) / gamma2
+        e = e_next
+        X = X + tau*D
+
+        if trust_region is not None:
+            if generic_vector_norm(X) > trust_region:
+                return _trust_tau(X, D, trust_region)
+
+        if (abs(beta) < eps) or (phi / b_norm <= tol):
+            # R = zeros(R)
+            return X
+
+        V_prev = V
+        V = P/beta
+        R = s**2*R - phi*c*V
+
+    return X

torchzero/utils/numberlist.py

@@ -129,4 +129,6 @@ class NumberList(list[int | float | Any]):
         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)
+    def clip(self, min=None, max=None):
         return self.zipmap_args(_clamp, min, max)

torchzero/utils/optimizer.py

@@ -1,3 +1,4 @@
+from abc import ABC, abstractmethod
 from collections.abc import Callable, Iterable, Mapping, MutableSequence, Sequence, MutableMapping
 from typing import Any, Literal, TypeVar, overload
 
@@ -132,65 +133,7 @@ def get_state_vals(state: Mapping[torch.Tensor, MutableMapping[str, Any]], param
     return values
 
 
-
-def loss_at_params(closure, params: Iterable[torch.Tensor],
-                   new_params: Sequence[torch.Tensor] | Any, backward: bool, restore=False):
-    params = TensorList(params)
-
-    old_params = params.clone() if restore else None
-
-    if isinstance(new_params, Sequence) and isinstance(new_params[0], torch.Tensor):
-        # when not restoring, copy new_params to params to avoid unexpected bugs due to shared storage
-        # when restoring params will be set back to old_params so its fine
-        if restore: params.set_(new_params)
-        else: params.copy_(new_params) # type:ignore
-
-    else:
-        new_params = totensor(new_params)
-        params.from_vec_(new_params)
-
-    if backward: loss = closure()
-    else: loss = closure(False)
-
-    if restore:
-        assert old_params is not None
-        params.set_(old_params)
-
-    return tofloat(loss)
-
-def loss_grad_at_params(closure, params: Iterable[torch.Tensor], new_params: Sequence[torch.Tensor], restore=False):
-    params = TensorList(params)
-    old_params = params.clone() if restore else None
-    loss = loss_at_params(closure, params, new_params, backward=True, restore=False)
-    grad = params.ensure_grad_().grad
-
-    if restore:
-        assert old_params is not None
-        params.set_(old_params)
-
-    return loss, grad
-
-def grad_at_params(closure, params: Iterable[torch.Tensor], new_params: Sequence[torch.Tensor], restore=False):
-    return loss_grad_at_params(closure=closure,params=params,new_params=new_params,restore=restore)[1]
-
-def loss_grad_vec_at_params(closure, params: Iterable[torch.Tensor], new_params: Any, restore=False):
-    params = TensorList(params)
-    old_params = params.clone() if restore else None
-    loss = loss_at_params(closure, params, new_params, backward=True, restore=False)
-    grad = params.ensure_grad_().grad.to_vec()
-
-    if restore:
-        assert old_params is not None
-        params.set_(old_params)
-
-    return loss, grad
-
-def grad_vec_at_params(closure, params: Iterable[torch.Tensor], new_params: Any, restore=False):
-    return loss_grad_vec_at_params(closure=closure,params=params,new_params=new_params,restore=restore)[1]
-
-
-
-class Optimizer(torch.optim.Optimizer):
+class Optimizer(torch.optim.Optimizer, ABC):
     """subclass of torch.optim.Optimizer with some helper methods for fast experimentation, it's not used anywhere in torchzero.
 
     Args:
@@ -251,21 +194,10 @@ class Optimizer(torch.optim.Optimizer):
 
         return get_state_vals(self.state, params, key, key2, *keys, init = init, cls = cls) # type:ignore[reportArgumentType]
 
-    def loss_at_params(self, closure, params: Sequence[torch.Tensor] | Any, backward: bool, restore=False):
-        return loss_at_params(closure=closure,params=self.get_params(),new_params=params,backward=backward,restore=restore)
-
-    def loss_grad_at_params(self, closure, params: Sequence[torch.Tensor] | Any, restore=False):
-        return loss_grad_at_params(closure=closure,params=self.get_params(),new_params=params,restore=restore)
-
-    def grad_at_params(self, closure, new_params: Sequence[torch.Tensor], restore=False):
-        return self.loss_grad_at_params(closure=closure,params=new_params,restore=restore)[1]
-
-    def loss_grad_vec_at_params(self, closure, params: Any, restore=False):
-        return loss_grad_vec_at_params(closure=closure,params=self.get_params(),new_params=params,restore=restore)
-
-    def grad_vec_at_params(self, closure, params: Any, restore=False):
-        return self.loss_grad_vec_at_params(closure=closure,params=params,restore=restore)[1]
 
+    # shut up pylance
+    @abstractmethod
+    def step(self, closure) -> Any: ... # pylint:disable=signature-differs # pyright:ignore[reportIncompatibleMethodOverride]
 
 def zero_grad_(params: Iterable[torch.Tensor], set_to_none):
     if set_to_none:
@@ -281,4 +213,53 @@ def zero_grad_(params: Iterable[torch.Tensor], set_to_none):
             else:
                 grad.requires_grad_(False)
 
-        torch._foreach_zero_(grads)
+        torch._foreach_zero_(grads)
+
+
+@overload
+def unpack_states(states: Sequence[MutableMapping[str, Any]], tensors: Sequence[torch.Tensor],
+                   key: str, *,
+                   must_exist: bool = False, init: Init = torch.zeros_like,
+                   cls: type[ListLike] = list) -> ListLike: ...
+@overload
+def unpack_states(states: Sequence[MutableMapping[str, Any]], tensors: Sequence[torch.Tensor],
+                   key: list[str] | tuple[str,...], *,
+                   must_exist: bool = False, init: Init | Sequence[Init] = torch.zeros_like,
+                   cls: type[ListLike] = list) -> list[ListLike]: ...
+@overload
+def unpack_states(states: Sequence[MutableMapping[str, Any]], tensors: Sequence[torch.Tensor],
+                   key: str,  key2: str, *keys: str,
+                   must_exist: bool = False, init: Init | Sequence[Init] = torch.zeros_like,
+                   cls: type[ListLike] = list) -> list[ListLike]: ...
+
+def unpack_states(states: Sequence[MutableMapping[str, Any]], tensors: Sequence[torch.Tensor],
+                   key: str | list[str] | tuple[str,...], key2: str | None = None,  *keys: str,
+                   must_exist: bool = False, init: Init | Sequence[Init] = torch.zeros_like,
+                   cls: type[ListLike] = list) -> ListLike | list[ListLike]:
+
+    # single key, return single cls
+    if isinstance(key, str) and key2 is None:
+        values = cls()
+        for i,s in enumerate(states):
+            if key not in s:
+                if must_exist: raise KeyError(f"Key {key} doesn't exist in state with keys {tuple(s.keys())}")
+                s[key] = _make_initial_state_value(tensors[i], init, i)
+            values.append(s[key])
+        return values
+
+    # multiple keys
+    k1 = (key,) if isinstance(key, str) else tuple(key)
+    k2 = () if key2 is None else (key2,)
+    keys = k1 + k2 + keys
+
+    values = [cls() for _ in keys]
+    for i,s in enumerate(states):
+        for k_i, key in enumerate(keys):
+            if key not in s:
+                if must_exist: raise KeyError(f"Key {key} doesn't exist in state with keys {tuple(s.keys())}")
+                k_init = init[k_i] if isinstance(init, (list,tuple)) else init
+                s[key] = _make_initial_state_value(tensors[i], k_init, i)
+            values[k_i].append(s[key])
+
+    return values
+

torchzero/utils/python_tools.py

@@ -1,7 +1,7 @@
 import functools
 import operator
-from typing import Any, TypeVar
-from collections.abc import Iterable, Callable
+from typing import Any, TypeVar, overload
+from collections.abc import Iterable, Callable, Mapping, MutableSequence
 from collections import UserDict
 
 
@@ -17,8 +17,8 @@ def flatten(iterable: Iterable) -> list[Any]:
     raise TypeError(f'passed object is not an iterable, {type(iterable) = }')
 
 X = TypeVar("X")
-# def reduce_dim[X](x:Iterable[Iterable[X]]) -> list[X]: # pylint:disable=E0602
-def reduce_dim(x:Iterable[Iterable[X]]) -> list[X]: # pylint:disable=E0602
+# def reduce_dim[X](x:Iterable[Iterable[X]]) -> list[X]:
+def reduce_dim(x:Iterable[Iterable[X]]) -> list[X]:
     """Reduces one level of nesting. Takes an iterable of iterables of X, and returns an iterable of X."""
     return functools.reduce(operator.iconcat, x, [])
 
@@ -31,6 +31,16 @@ def generic_eq(x: int | float | Iterable[int | float], y: int | float | Iterable
         return all(i==y for i in x)
     return all(i==j for i,j in zip(x,y))
 
+def generic_ne(x: int | float | Iterable[int | float], y: int | float | Iterable[int | float]) -> bool:
+    """generic not equals function that supports scalars and lists of numbers. Faster than not generic_eq"""
+    if isinstance(x, (int,float)):
+        if isinstance(y, (int,float)): return x!=y
+        return any(i!=x for i in y)
+    if isinstance(y, (int,float)):
+        return any(i!=y for i in x)
+    return any(i!=j for i,j in zip(x,y))
+
+
 def zipmap(self, fn: Callable, other: Any | list | tuple, *args, **kwargs):
     """If `other` is list/tuple, applies `fn` to self zipped with `other`.
     Otherwise applies `fn` to this sequence and `other`.
@@ -38,3 +48,16 @@ def zipmap(self, fn: Callable, other: Any | list | tuple, *args, **kwargs):
     if isinstance(other, (list, tuple)): return self.__class__(fn(i, j, *args, **kwargs) for i, j in zip(self, other))
     return self.__class__(fn(i, other, *args, **kwargs) for i in self)
 
+ListLike = TypeVar('ListLike', bound=MutableSequence)
+@overload
+def unpack_dicts(dicts: Iterable[Mapping[str, Any]], key:str, *, cls:type[ListLike]=list) -> ListLike: ...
+@overload
+def unpack_dicts(dicts: Iterable[Mapping[str, Any]], key:str, key2: str, *keys:str, cls:type[ListLike]=list) -> list[ListLike]: ...
+def unpack_dicts(dicts: Iterable[Mapping[str, Any]], key:str, key2: str | None = None, *keys:str, cls:type[ListLike]=list) -> ListLike | list[ListLike]:
+    k1 = (key,) if isinstance(key, str) else tuple(key)
+    k2 = () if key2 is None else (key2,)
+    keys = k1 + k2 + keys
+
+    values = [cls(s[k] for s in dicts) for k in keys] # pyright:ignore[reportCallIssue]
+    if len(values) == 1: return values[0]
+    return values