torchzero 0.4.2__tar.gz → 0.4.4__tar.gz

{torchzero-0.4.2 → torchzero-0.4.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchzero
-Version: 0.4.2
+Version: 0.4.4
 Summary: Modular optimization library for PyTorch.
 Author-email: Ivan Nikishev <nkshv2@gmail.com>
 Project-URL: Homepage, https://github.com/inikishev/torchzero

{torchzero-0.4.2 → torchzero-0.4.4}/pyproject.toml

@@ -13,7 +13,7 @@ build-backend = "setuptools.build_meta"
 name = "torchzero"
 description = "Modular optimization library for PyTorch."
 
-version = "0.4.2"
+version = "0.4.4"
 dependencies = [
   "torch",
   "numpy",

{torchzero-0.4.2 → torchzero-0.4.4}/tests/test_identical.py

@@ -105,7 +105,7 @@ def test_adam(amsgrad):
     tz_fn_ops = lambda p: tz.Optimizer(
         p,
         tz.m.DivModules(
-            tz.m.EMA(0.9, debiased=True),
+            tz.m.EMA(0.9, debias=True),
             [tz.m.SqrtEMASquared(0.999, debiased=True, amsgrad=amsgrad), tz.m.Add(1e-8)]
         ))
     tz_fn_ops2 = lambda p: tz.Optimizer(

{torchzero-0.4.2 → torchzero-0.4.4}/tests/test_opts.py

@@ -727,8 +727,8 @@ Adam = Run(
 )
 # ------------------------------ optimizers/soap ----------------------------- #
 SOAP = Run(
-    func_opt=lambda p: tz.Optimizer(p, tz.m.SOAP(), tz.m.LR(0.4)),
-    sphere_opt=lambda p: tz.Optimizer(p, tz.m.SOAP(precond_freq=1), tz.m.LR(1)),
+    func_opt=lambda p: tz.Optimizer(p, tz.m.SOAP(merge_small=True), tz.m.LR(0.4)),
+    sphere_opt=lambda p: tz.Optimizer(p, tz.m.SOAP(precond_freq=1, merge_small=True), tz.m.LR(1)),
     needs_closure=False,
     # merge and unmerge lrs are very different so need to test convergence separately somewhere
     func='rosen', steps=50, loss=4, merge_invariant=False,

{torchzero-0.4.2 → torchzero-0.4.4}/torchzero/_minimize/methods.py

@@ -14,82 +14,87 @@ from ..utils import tofloat
 
 
 def _get_method_from_str(method: str) -> list[Module]:
-    method = ''.join(c for c in method.lower().strip() if c.isalnum())
+    stripped = ''.join(c for c in method.lower().strip() if c.isalnum())
 
-    if method == "bfgs":
+    if stripped == "bfgs":
         return [m.RestartOnStuck(m.BFGS()), m.Backtracking()]
 
-    if method == "lbfgs":
+    if stripped == "lbfgs":
         return [m.LBFGS(100), m.Backtracking()]
 
-    if method == "newton":
+    if stripped == "newton":
         return [m.Newton(), m.Backtracking()]
 
-    if method == "sfn":
+    if stripped == "sfn":
         return [m.Newton(eigval_fn=lambda x: x.abs().clip(min=1e-10)), m.Backtracking()]
 
-    if method == "inm":
+    if stripped == "inm":
         return [m.ImprovedNewton(), m.Backtracking()]
 
-    if method == 'crn':
+    if stripped == 'crn':
         return [m.CubicRegularization(m.Newton())]
 
-    if method == "commondirections":
+    if stripped == "commondirections":
         return [m.SubspaceNewton(sketch_type='common_directions'), m.Backtracking()]
 
-    if method == "trust":
+    if stripped == "trust":
         return [m.LevenbergMarquardt(m.Newton())]
 
-    if method == "trustexact":
-        return [m.TrustCG(m.Newton())]
-
-    if method == "dogleg":
+    if stripped == "dogleg":
         return [m.Dogleg(m.Newton())]
 
-    if method == "trustbfgs":
-        return [m.LevenbergMarquardt(m.BFGS())]
+    if stripped == "trustbfgs":
+        return [m.RestartOnStuck(m.LevenbergMarquardt(m.BFGS()))]
 
-    if method == "trustsr1":
-        return [m.LevenbergMarquardt(m.SR1())]
+    if stripped == "trustsr1":
+        return [m.RestartOnStuck(m.LevenbergMarquardt(m.SR1()))]
 
-    if method == "newtoncg":
+    if stripped == "newtoncg":
         return [m.NewtonCG(), m.Backtracking()]
 
-    if method == "tn":
+    if stripped == "tn":
         return [m.NewtonCG(maxiter=10), m.Backtracking()]
 
-    if method == "trustncg":
+    if stripped == "trustncg":
         return [m.NewtonCGSteihaug()]
 
-    if method == "gd":
+    if stripped == "gd":
         return [m.Backtracking()]
 
-    if method == "cg":
+    if stripped == "cg":
         return [m.FletcherReeves(), m.StrongWolfe(c2=0.1, fallback=True)]
 
-    if method == "bb":
+    if stripped in ("shor", "shorr"):
+        return [m.ShorR(), m.StrongWolfe(c2=0.1, fallback=True)]
+
+    if stripped == "pgm":
+        return [m.ProjectedGradientMethod(), m.StrongWolfe(c2=0.1, fallback=True)]
+
+    if stripped == "bb":
         return [m.RestartOnStuck(m.BarzilaiBorwein())]
 
-    if method == "bbstab":
+    if stripped == "bbstab":
         return [m.BBStab()]
 
-    if method == "adgd":
+    if stripped == "adgd":
         return [m.AdGD()]
 
-    if method in ("gn", "gaussnewton"):
+    if stripped in ("bd", "bolddriver"):
+        return [m.BoldDriver()]
+
+    if stripped in ("gn", "gaussnewton"):
         return [m.GaussNewton(), m.Backtracking()]
 
-    if method == "rprop":
+    if stripped == "rprop":
         return [m.Rprop(alpha=1e-3)]
 
-    if method == "lm":
+    if stripped == "lm":
         return [m.LevenbergMarquardt(m.GaussNewton())]
 
-    if method == "mlm":
+    if stripped == "mlm":
         return [m.LevenbergMarquardt(m.GaussNewton(), y=1)]
 
-    if method == "cd":
+    if stripped == "cd":
         return [m.CD(), m.ScipyMinimizeScalar(maxiter=8)]
 
-
-    raise NotImplementedError(method)
+    raise NotImplementedError(stripped)

{torchzero-0.4.2 → torchzero-0.4.4}/torchzero/linalg/eigh.py

@@ -285,8 +285,8 @@ def rank1_eigh(v: torch.Tensor):
     vv = v.dot(v)
     norm = vv.sqrt().clip(min=torch.finfo(vv.dtype).tiny * 2)
 
-    L = vv.unsqueeze(0) # (rank, )
-    Q = v.unsqueeze(-1) / norm # (m, rank)
+    L = vv.unsqueeze(0) # (1, )
+    Q = v.unsqueeze(-1) / norm # (m, 1)
 
     return L, Q
 

{torchzero-0.4.2 → torchzero-0.4.4}/torchzero/linalg/torch_linalg.py

@@ -46,7 +46,7 @@ def eigh(A: torch.Tensor, UPLO="L", retry_float64:bool=False) -> tuple[torch.Ten
     try:
         return torch.linalg.eigh(A, UPLO=UPLO) # pylint:disable=not-callable
 
-    except torch.linalg.LinAlgError as e:
+    except (torch.linalg.LinAlgError, RuntimeError) as e:
         if not retry_float64: raise e
         dtype = A.dtype
         if dtype == torch.float64: raise e

{torchzero-0.4.2 → torchzero-0.4.4}/torchzero/modules/adaptive/ggt.py

@@ -130,7 +130,7 @@ class GGT(TensorTransform):
         step = state.get('step', 0)
         state['step'] = step + 1
 
-        if step % update_freq == 0 :
+        if step % update_freq == 0:
 
             # compute new factors
             L = state.get("L", None)

{torchzero-0.4.2 → torchzero-0.4.4}/torchzero/modules/adaptive/sam.py

@@ -1,7 +1,10 @@
+from collections.abc import Mapping, Sequence
 from contextlib import nullcontext
+from typing import Any
 import torch
+
 from ...utils import TensorList, NumberList, unpack_dicts, unpack_states
-from ...core import Transform
+from ...core import Transform, Objective
 
 
 class SAM(Transform):
@@ -126,6 +129,8 @@ class SAM(Transform):
 
         objective.closure = sam_closure
 
+    def apply_states(self, objective: Objective, states: list[dict[str, Any]], settings: Sequence[Mapping[str, Any]]) -> Objective:
+        return objective
 # different class because defaults for SAM are bad for ASAM
 class ASAM(SAM):
     """Adaptive Sharpness-Aware Minimization from https://arxiv.org/pdf/2102.11600#page=6.52

{torchzero-0.4.2 → torchzero-0.4.4}/torchzero/modules/adaptive/shampoo.py

@@ -31,7 +31,7 @@ def update_shampoo_preconditioner_(
             if reg != 0:
                 accumulator = accumulator + torch.eye(accumulator.size(0), device=accumulator.device, dtype=accumulator.dtype).mul_(reg)
 
-            if matrix_power is None: matrix_power = -1 / max(grad.ndim, 2)
+            if matrix_power is None: matrix_power = -1 / max(grad.ndim * 2, 2)
             set_storage_(preconditioner, _matrix_power(accumulator, matrix_power, method=matrix_power_method))
 
 def apply_shampoo_preconditioner(

{torchzero-0.4.2 → torchzero-0.4.4}/torchzero/modules/adaptive/soap.py

@@ -51,6 +51,7 @@ def project_back(tensor: torch.Tensor, Q: list[torch.Tensor| None]):
     return tensor
 
 # function from https://github.com/nikhilvyas/SOAP/blob/main/soap.py
+# this is only used once per accumulator to initialize it
 @torch.no_grad
 def get_orthogonal_matrix(mats: list[torch.Tensor | None]):
     """
@@ -64,7 +65,19 @@ def get_orthogonal_matrix(mats: list[torch.Tensor | None]):
             final.append(None)
             continue
 
-        _, Q = torch_linalg.eigh(M + 1e-30 * torch.eye(M.shape[0], device=M.device), retry_float64=True)
+        if not torch.isfinite(M).all():
+            raise RuntimeError(f"Initial gradient for parameter {M.shape} has non-finite values.")
+
+        M_f64 = M.to(torch.float64) + 1e-30 * torch.eye(M.shape[0], device=M.device, dtype=torch.float64)
+        try:
+            _, Q_f64 = torch_linalg.eigh(M_f64)
+        except RuntimeError as e:
+            if M_f64.is_cpu: raise e
+            M_f64 = M_f64.cpu()
+            _, Q_f64 = torch_linalg.eigh(M_f64) # apparently there is a bug in CUDA eigh
+            Q_f64 = Q_f64.to(M.device)
+
+        Q = Q_f64.to(M.dtype)
 
         Q = torch.flip(Q, [1])
         final.append(Q)
@@ -156,7 +169,7 @@ class SOAP(TensorTransform):
         beta2: float = 0.95,
         shampoo_beta: float | None = 0.95,
         precond_freq: int = 10,
-        merge_small: bool = True,
+        merge_small: bool = False,
         max_dim: int = 4096,
         precondition_1d: bool = True,
         eps: float = 1e-8,

{torchzero-0.4.2 → torchzero-0.4.4}/torchzero/modules/basis/ggt_basis.py

@@ -111,7 +111,7 @@ class GGTBasis(TensorTransform):
         inner: Chainable | None = None,
     ):
         defaults = locals().copy()
-        del defaults['self'], defaults['inner']
+        del defaults['self'], defaults['inner'], defaults["basis_opt"]
 
         super().__init__(defaults, concat_params=True, inner=inner)
         self.set_child("basis_opt", basis_opt)

{torchzero-0.4.2 → torchzero-0.4.4}/torchzero/modules/misc/multistep.py

@@ -154,8 +154,7 @@ class Online(Module):
         closure = objective.closure
         if closure is None: raise ValueError("Closure must be passed for Online")
 
-        step = self.global_state.get('step', 0) + 1
-        self.global_state['step'] = step
+        step = self.increment_counter("step", start = 0)
 
         params = TensorList(objective.params)
         p_cur = params.clone()
@@ -165,7 +164,7 @@ class Online(Module):
         var_c = objective.clone(clone_updates=False)
 
         # on 1st step just step and store previous params
-        if step == 1:
+        if step == 0:
             p_prev.copy_(params)
 
             module.update(var_c)

{torchzero-0.4.2 → torchzero-0.4.4}/torchzero/modules/misc/split.py

@@ -53,11 +53,11 @@ _SingleFilter = Callable[[torch.Tensor], bool] | torch.Tensor | Iterable[torch.T
 Filter = _SingleFilter | Iterable[_SingleFilter]
 
 def _make_filter(filter: Filter):
-    if callable(filter): return filter
     if isinstance(filter, torch.Tensor):
         return lambda x: x is filter
     if isinstance(filter, torch.nn.Module):
         return _make_filter(filter.parameters())
+    if callable(filter): return filter
 
     # iterable
     filters = [_make_filter(f) for f in filter]

{torchzero-0.4.2 → torchzero-0.4.4}/torchzero/modules/momentum/momentum.py

@@ -6,7 +6,7 @@ import torch
 
 from ...core import  TensorTransform
 from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
-from ..opt_utils import debias, ema_
+from ..opt_utils import debias as _debias, ema_
 
 
 class EMA(TensorTransform):
@@ -15,13 +15,13 @@ class EMA(TensorTransform):
     Args:
         momentum (float, optional): momentum (beta). Defaults to 0.9.
         dampening (float, optional): momentum dampening. Defaults to 0.
-        debiased (bool, optional): whether to debias the EMA like in Adam. Defaults to False.
+        debias (bool, optional): whether to debias the EMA like in Adam. Defaults to False.
         lerp (bool, optional): whether to use linear interpolation. Defaults to True.
         ema_init (str, optional): initial values for the EMA, "zeros" or "update".
         target (Target, optional): target to apply EMA to. Defaults to 'update'.
     """
-    def __init__(self, momentum:float=0.9, dampening:float=0, debiased: bool = False, lerp=True, ema_init: Literal['zeros', 'update'] = 'zeros'):
-        defaults = dict(momentum=momentum,dampening=dampening,debiased=debiased,lerp=lerp,ema_init=ema_init)
+    def __init__(self, momentum:float=0.9, dampening:float=0, debias: bool = False, lerp=True, ema_init: Literal['zeros', 'update'] = 'zeros'):
+        defaults = dict(momentum=momentum,dampening=dampening,debias=debias,lerp=lerp,ema_init=ema_init)
         super().__init__(defaults, uses_grad=False)
 
         self.add_projected_keys("grad", "exp_avg")
@@ -30,7 +30,7 @@ class EMA(TensorTransform):
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         step = self.global_state['step'] = self.global_state.get('step', 0) + 1
 
-        debiased, lerp, ema_init = itemgetter('debiased','lerp','ema_init')(settings[0])
+        debias, lerp, ema_init = itemgetter('debias','lerp','ema_init')(settings[0])
 
         exp_avg = unpack_states(states, tensors, 'exp_avg',
                                 init=torch.zeros_like if ema_init=='zeros' else tensors, cls=TensorList)
@@ -38,7 +38,7 @@ class EMA(TensorTransform):
 
         exp_avg = ema_(TensorList(tensors), exp_avg_=exp_avg,beta=momentum,dampening=dampening,lerp=lerp)
 
-        if debiased: return debias(exp_avg, step=step, beta1=momentum, alpha=1, inplace=False)
+        if debias: return _debias(exp_avg, step=step, beta1=momentum, alpha=1, inplace=False)
         else: return exp_avg.clone() # this has exp_avg storage so needs to be cloned
 
 
@@ -49,14 +49,14 @@ class HeavyBall(EMA):
     Args:
         momentum (float, optional): momentum (beta). Defaults to 0.9.
         dampening (float, optional): momentum dampening. Defaults to 0.
-        debiased (bool, optional): whether to debias the EMA like in Adam. Defaults to False.
+        debias (bool, optional): whether to debias the EMA like in Adam. Defaults to False.
         lerp (bool, optional):
             whether to use linear interpolation, if True, this becomes exponential moving average. Defaults to False.
         ema_init (str, optional): initial values for the EMA, "zeros" or "update".
         target (Target, optional): target to apply EMA to. Defaults to 'update'.
     """
-    def __init__(self, momentum:float=0.9, dampening:float=0, debiased: bool = False, lerp=False, ema_init: Literal['zeros', 'update'] = 'update'):
-        super().__init__(momentum=momentum, dampening=dampening, debiased=debiased, lerp=lerp, ema_init=ema_init)
+    def __init__(self, momentum:float=0.9, dampening:float=0, debias: bool = False, lerp=False, ema_init: Literal['zeros', 'update'] = 'update'):
+        super().__init__(momentum=momentum, dampening=dampening, debias=debias, lerp=lerp, ema_init=ema_init)
 
 def nag_(
     tensors_: TensorList,