torchzero 0.4.3__tar.gz → 0.4.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchzero
-Version: 0.4.3
+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.3 → 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.3"
+version = "0.4.4"
 dependencies = [
   "torch",
   "numpy",

{torchzero-0.4.3 → 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.3 → 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.3 → 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.3 → 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.3 → 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.3 → 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.3 → 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.3 → 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.3 → 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.3 → 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.3 → torchzero-0.4.4}/torchzero/modules/second_order/newton.py

@@ -44,7 +44,14 @@ def _newton_update_state_(
 
     # if any args require eigendecomp, we don't need H or H_inv, we store factors
     if any(i is not None for i in [eigval_fn, eigv_tol, truncate]):
-        L, Q = torch_linalg.eigh(H, retry_float64=True)
+        try:
+            state.pop("H", None)
+            L, Q = torch_linalg.eigh(H, retry_float64=True)
+        except torch.linalg.LinAlgError:
+            state.pop("L",None); state.pop("Q",None)
+            state["H"] = H
+            return
+
         if eigval_fn is not None: L = eigval_fn(L)
         L, Q = regularize_eigh(L, Q, truncate=truncate, tol=eigv_tol)
         state["L"] = L

{torchzero-0.4.3 → torchzero-0.4.4}/torchzero/modules/step_size/__init__.py

@@ -1,2 +1,2 @@
 from .lr import LR, StepSize, Warmup, WarmupNormClip, RandomStepSize
-from .adaptive import PolyakStepSize, BarzilaiBorwein, BBStab, AdGD
+from .adaptive import PolyakStepSize, BarzilaiBorwein, BBStab, AdGD, BoldDriver