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

tests/test_opts.py

@@ -745,7 +745,7 @@ SSVM = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.SSVM(1), tz.m.StrongWolfe()),
     sphere_opt=lambda p: tz.Modular(p, tz.m.SSVM(1), tz.m.StrongWolfe()),
     needs_closure=True,
-    func='rosen', steps=50, loss=1e-12, merge_invariant=True,
+    func='rosen', steps=50, loss=1e-10, merge_invariant=True,
     sphere_steps=10, sphere_loss=0,
 )
 

tests/test_tensorlist.py

@@ -835,7 +835,7 @@ def test_global_reductions(simple_tl: TensorList, global_method, vec_equiv_metho
     expected = vec_equiv_func()
 
     if isinstance(result, bool): assert result == expected
-    else: assert torch.allclose(result, expected), f"Tensors not close: {result = }, {expected = }"
+    else: assert torch.allclose(result, expected, atol=1e-4), f"Tensors not close: {result = }, {expected = }"
 
 
 def test_global_vector_norm(simple_tl: TensorList):

torchzero/core/preconditioner.py

@@ -38,17 +38,18 @@ class Preconditioner(Transform):
 
 
     def _tensor_wise_transform(self, tensors:list[torch.Tensor], params:list[torch.Tensor], grads:list[torch.Tensor] | None, vars:Vars) -> list[torch.Tensor]:
-        step = self.global_state.get('step', 0)
+        step = self.global_state.get('__step', 0)
         states = [self.state[p] for p in params]
         settings = [self.settings[p] for p in params]
         global_settings = settings[0]
         update_freq = global_settings['__update_freq']
 
         scale_first = global_settings['__scale_first']
-        scale_factor = 0
+        scale_factor = 1
         if scale_first and step == 0:
             # initial step size guess from pytorch LBFGS
-            scale_factor = TensorList(tensors).abs().sum()
+            scale_factor = 1 / TensorList(tensors).abs().global_sum().clip(min=1)
+            scale_factor = scale_factor.clip(min=torch.finfo(tensors[0].dtype).eps)
 
         # update preconditioner
         if step % update_freq == 0:
@@ -63,13 +64,13 @@ class Preconditioner(Transform):
 
         # scale initial step, when preconditioner might not have been applied
         if scale_first and step == 0:
-            torch._foreach_div_(tensors, scale_factor)
+            torch._foreach_mul_(tensors, scale_factor)
 
-        self.global_state['step'] = step + 1
+        self.global_state['__step'] = step + 1
         return tensors
 
     def _concat_transform(self, tensors:list[torch.Tensor], params:list[torch.Tensor], grads:list[torch.Tensor] | None, vars:Vars) -> list[torch.Tensor]:
-        step = self.global_state.get('step', 0)
+        step = self.global_state.get('__step', 0)
         tensors_vec = torch.cat([t.ravel() for t in tensors])
         params_vec = torch.cat([p.ravel() for p in params])
         grads_vec = [torch.cat([g.ravel() for g in grads])] if grads is not None else None
@@ -80,10 +81,11 @@ class Preconditioner(Transform):
         update_freq = global_settings['__update_freq']
 
         scale_first = global_settings['__scale_first']
-        scale_factor = 0
+        scale_factor = 1
         if scale_first and step == 0:
             # initial step size guess from pytorch LBFGS
-            scale_factor = tensors_vec.abs().sum()
+            scale_factor = 1 / tensors_vec.abs().sum().clip(min=1)
+            scale_factor = scale_factor.clip(min=torch.finfo(tensors_vec.dtype).eps)
 
         # update preconditioner
         if step % update_freq == 0:
@@ -99,11 +101,10 @@ class Preconditioner(Transform):
 
         # scale initial step, when preconditioner might not have been applied
         if scale_first and step == 0:
-            if scale_factor >= torch.finfo(tensors_vec.dtype).eps:
-                tensors_vec /= scale_factor
+            tensors_vec *= scale_factor
 
         tensors = vec_to_tensors(vec=tensors_vec, reference=tensors)
-        self.global_state['step'] = step + 1
+        self.global_state['__step'] = step + 1
         return tensors
 
     @torch.no_grad

torchzero/modules/experimental/__init__.py

@@ -3,7 +3,7 @@ from .adadam import Adadam
 from .adamY import AdamY
 from .adasoap import AdaSOAP
 from .curveball import CurveBall
-from .dsoap import DSOAP
+from .soapy import SOAPY
 from .gradmin import GradMin
 from .reduce_outward_lr import ReduceOutwardLR
 from .spectral import SpectralPreconditioner
@@ -11,4 +11,5 @@ from .subspace_preconditioners import (
     HistorySubspacePreconditioning,
     RandomSubspacePreconditioning,
 )
-from .tropical_newton import TropicalNewton
+from .tropical_newton import TropicalNewton
+from .newton_solver import NewtonSolver

torchzero/modules/experimental/absoap.py

@@ -140,11 +140,17 @@ Source=Literal['p','g','s','y', 'gy', 'sy', 'sn', 'yn', 'gys', 'sys','sn', 'yn']
 class ABSOAP(Transform):
     """SOAP but with two extra letters included in its name in order to improve converence
 
+    so what you can do is choose what goes into what ,and that is supposed to be good.
+
     new args
 
     scale by s whether to scale gradient differences by parameter differences
 
     y_to_ema2 whether to use gradient differences for exponential moving average too
+
+    okay I changed these args into another ones
+
+    BASICALLY THIS IS FOR MY EXPERIMENTS
     """
     def __init__(
         self,
@@ -213,7 +219,7 @@ class ABSOAP(Transform):
             if 'g_prev' not in state:
                 state['p_prev'] = p.clone()
                 state['g_prev'] = t.clone()
-                updates.append(tensors[i].sign())
+                updates.append(tensors[i].clip(-0.1,0.1))
                 continue
 
             p_prev = state['p_prev']
@@ -285,7 +291,7 @@ class ABSOAP(Transform):
                     state['Q'] = get_orthogonal_matrix(state['GG'])
 
                 state['step'] = 0
-                updates.append(tensors[i].sign())
+                updates.append(tensors[i].clip(-0.1,0.1))
                 continue  # skip 1st step as in https://github.com/nikhilvyas/SOAP/blob/main/soap.py ?
                 # I use sign instead as to not mess up with next modules. 1st Adam step is always sign anyway.
 

torchzero/modules/experimental/adadam.py

@@ -50,7 +50,7 @@ def adadam_(
     return None
 
 class Adadam(Module):
-    """Adam with a diagonally preconditioned preconditioner and a graceful name."""
+    """Adam with a diagonally preconditioned preconditioner."""
     def __init__(
         self,
         beta1: float = 0.9,

torchzero/modules/experimental/adamY.py

@@ -37,7 +37,7 @@ def adamy_(
         p_prev.copy_(p)
         g_prev.copy_(g)
 
-        update = g.sign().lazy_mul_(alpha*0.1)
+        update = g.clip(-0.1,0.1).lazy_mul_(alpha)
         if params_ is None: return update
         params_.sub_(update)
         return None

torchzero/modules/experimental/adasoap.py

@@ -218,9 +218,9 @@ class AdaSOAP(Transform):
                     state['Q'] = get_orthogonal_matrix(GG_precond)
 
                 state['step'] = 0
-                updates.append(tensors[i].sign())
+                updates.append(tensors[i].clip(-0.1,0.1))
                 continue  # skip 1st step as in https://github.com/nikhilvyas/SOAP/blob/main/soap.py ?
-                # I use sign instead as to not mess up with next modules. 1st Adam step is always sign anyway.
+                # that can mess with other modules scaling
 
             # Projecting gradients to the eigenbases of Shampoo's preconditioner
             # i.e. projecting to the eigenbases of matrices in state['GG']

torchzero/modules/experimental/algebraic_newton.py

@@ -71,7 +71,7 @@ def tikhonov(H: torch.Tensor, reg: float, algebra: ta.Algebra = ta.TropicalSemir
 
 
 class AlgebraicNewton(Module):
-    """newton in other algebras, not practical because solving linear system is very hard."""
+    """newton in other algebras, not that it works."""
     def __init__(
         self,
         reg: float | None = None,

torchzero/modules/experimental/curveball.py

@@ -13,7 +13,7 @@ def curveball(
     momentum: float | NumberList,
     precond_lr: float | NumberList,
 ):
-    """returns z_, clone it!!!"""
+    """returns z_, clone it!!! (no just negate it)"""
     delta = Hz + tensors
     z_.mul_(momentum).sub_(delta.mul_(precond_lr)) # z ← ρz − βΔ
     return z_

torchzero/modules/experimental/gradmin.py

@@ -14,7 +14,7 @@ from ..smoothing.gaussian import Reformulation
 
 
 class GradMin(Reformulation):
-    """Reformulates the objective to minimize sum of gradient magnitudes via autograd.
+    """Reformulates the objective to minimize sum of gradient magnitudes via autograd. This is not expected to be practical.
 
     Args:
         loss_term (float, optional): adds loss value times this to sum of gradient magnitudes. Defaults to 1.

torchzero/modules/experimental/newton_solver.py

@@ -0,0 +1,88 @@
+from collections.abc import Callable, Iterable
+from typing import Any, Literal, overload
+
+import torch
+
+from ...core import Chainable, Module, apply, Modular
+from ...utils import TensorList, as_tensorlist
+from ...utils.derivatives import hvp
+from ..quasi_newton import LBFGS
+
+class NewtonSolver(Module):
+    """Matrix free newton via with any custom solver (usually it is better to just use NewtonCG or NystromPCG is even better)"""
+    def __init__(
+        self,
+        solver: Callable[[list[torch.Tensor]], Any] = lambda p: Modular(p, LBFGS()),
+        maxiter=None,
+        tol=1e-3,
+        reg: float = 0,
+        warm_start=True,
+        inner: Chainable | None = None,
+    ):
+        defaults = dict(tol=tol, maxiter=maxiter, reg=reg, warm_start=warm_start, solver=solver)
+        super().__init__(defaults,)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    @torch.no_grad
+    def step(self, vars):
+        params = TensorList(vars.params)
+        closure = vars.closure
+        if closure is None: raise RuntimeError('NewtonCG requires closure')
+
+        settings = self.settings[params[0]]
+        solver_cls = settings['solver']
+        maxiter = settings['maxiter']
+        tol = settings['tol']
+        reg = settings['reg']
+        warm_start = settings['warm_start']
+
+        # ---------------------- Hessian vector product function --------------------- #
+        grad = vars.get_grad(create_graph=True)
+
+        def H_mm(x):
+            with torch.enable_grad():
+                Hvp = TensorList(hvp(params, grad, x, create_graph=True))
+                if reg != 0: Hvp = Hvp + (x*reg)
+                return Hvp
+
+        # -------------------------------- inner step -------------------------------- #
+        b = as_tensorlist(grad)
+        if 'inner' in self.children:
+            b = as_tensorlist(apply(self.children['inner'], [g.clone() for g in grad], params=params, grads=grad, vars=vars))
+
+        # ---------------------------------- run cg ---------------------------------- #
+        x0 = None
+        if warm_start: x0 = self.get_state('prev_x', params=params, cls=TensorList) # initialized to 0 which is default anyway
+        if x0 is None: x = b.zeros_like().requires_grad_(True)
+        else: x = x0.clone().requires_grad_(True)
+
+        solver = solver_cls(x)
+        def lstsq_closure(backward=True):
+            Hx = H_mm(x)
+            loss = (Hx-b).pow(2).global_mean()
+            if backward:
+                solver.zero_grad()
+                loss.backward(inputs=x)
+            return loss
+
+        if maxiter is None: maxiter = b.global_numel()
+        loss = None
+        initial_loss = lstsq_closure(False)
+        if initial_loss > tol:
+            for i in range(maxiter):
+                loss = solver.step(lstsq_closure)
+                assert loss is not None
+                if min(loss, loss/initial_loss) < tol: break
+
+        print(f'{loss = }')
+
+        if warm_start:
+            assert x0 is not None
+            x0.copy_(x)
+
+        vars.update = x.detach()
+        return vars
+
+

torchzero/modules/experimental/{dsoap.py → soapy.py}

@@ -3,7 +3,7 @@ from operator import itemgetter
 import torch
 
 from ...core import Chainable, Transform, apply
-from ...modules.optimizers.shampoo import _merge_small_dims, _unmerge_small_dims
+from ..optimizers.shampoo import _merge_small_dims, _unmerge_small_dims
 
 @torch.no_grad
 def update_soap_covariances_(
@@ -135,7 +135,7 @@ def get_orthogonal_matrix_QR(exp_avg_sq: torch.Tensor, GG: list[torch.Tensor | N
 
     return final, exp_avg_sq
 
-class DSOAP(Transform):
+class SOAPY(Transform):
     """SOAP but uses scaled gradient differences
 
     new args
@@ -195,7 +195,7 @@ class DSOAP(Transform):
             if 'g_prev' not in state:
                 state['p_prev'] = p.clone()
                 state['g_prev'] = t.clone()
-                updates.append(tensors[i].sign())
+                updates.append(tensors[i].clip(-0.1,0.1))
                 continue
 
             p_prev = state['p_prev']
@@ -228,7 +228,7 @@ class DSOAP(Transform):
                     state['Q'] = get_orthogonal_matrix(state['GG'])
 
                 state['step'] = 0
-                updates.append(tensors[i].sign())
+                updates.append(tensors[i].clip(-0.1,0.1))
                 continue  # skip 1st step as in https://github.com/nikhilvyas/SOAP/blob/main/soap.py ?
                 # I use sign instead as to not mess up with next modules. 1st Adam step is always sign anyway.
 

torchzero/modules/experimental/spectral.py

@@ -194,8 +194,10 @@ class SpectralPreconditioner(TensorwisePreconditioner):
         order (int, optional):
             whitening order, 1 approximates FIM (maybe), 2 - hessian (maybe), 3+ - god knows what.
         solver (str, optional): what to use for whitening. Defaults to 'svd'.
-        U_beta (float | None, optional): beta for U (probably a bad idea). Defaults to None.
-        S_beta (float | None, optional): beta for S (probably a bad idea). Defaults to None.
+        A_beta (float | None, optional):
+            beta for U (in SVD and other letters in other solvers) (probably a bad idea). Defaults to None.
+        B_beta (float | None, optional):
+            beta for S (in SVD and other letters in other solvers) (probably a bad idea). Defaults to None.
         interval (int, optional): How often to update history. Defaults to 1 (every step).
         concat_params (bool, optional):
             whether to apply preconditioning to each tensor (False, default) or to all tensors concatenated into a vector (True). Latter will be slower but captures interactions between layers. Defaults to True.
@@ -275,7 +277,7 @@ class SpectralPreconditioner(TensorwisePreconditioner):
         A = state.get('A', None)
         if A is None:
             # make a conservative step to avoid issues due to different GD scaling
-            return tensor.div_(max(1, tensor.abs().sum())) # pyright:ignore[reportArgumentType]
+            return tensor.clip_(-0.1, 0.1) # pyright:ignore[reportArgumentType]
 
         B = state['B']
         update = solver.apply(tensor.view(-1), A, B).view_as(tensor)

torchzero/modules/experimental/structured_newton.py

@@ -0,0 +1,111 @@
+# idea https://arxiv.org/pdf/2212.09841
+import warnings
+from collections.abc import Callable
+from functools import partial
+from typing import Literal
+
+import torch
+
+from ...core import Chainable, Module, apply
+from ...utils import TensorList, vec_to_tensors
+from ...utils.derivatives import (
+    hessian_list_to_mat,
+    hessian_mat,
+    hvp,
+    hvp_fd_central,
+    hvp_fd_forward,
+    jacobian_and_hessian_wrt,
+)
+
+
+class StructuredNewton(Module):
+    """TODO
+    Args:
+        structure (str, optional): structure.
+        reg (float, optional): tikhonov regularizer value. Defaults to 1e-6.
+        hvp_method (str):
+            how to calculate hvp_method. Defaults to "autograd".
+        inner (Chainable | None, optional): inner modules. Defaults to None.
+
+    """
+    def __init__(
+        self,
+        structure: Literal[
+            "diagonal",
+            "diagonal1",
+            "diagonal_abs",
+            "tridiagonal",
+            "circulant",
+            "toeplitz",
+            "toeplitz_like",
+            "hankel",
+            "rank1",
+            "rank2", # any rank
+        ]
+        | str = "diagonal",
+        reg: float = 1e-6,
+        hvp_method: Literal["autograd", "forward", "central"] = "autograd",
+        h: float = 1e-3,
+        inner: Chainable | None = None,
+    ):
+        defaults = dict(reg=reg, hvp_method=hvp_method, structure=structure, h=h)
+        super().__init__(defaults)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    @torch.no_grad
+    def step(self, vars):
+        params = TensorList(vars.params)
+        closure = vars.closure
+        if closure is None: raise RuntimeError('NewtonCG requires closure')
+
+        settings = self.settings[params[0]]
+        reg = settings['reg']
+        hvp_method = settings['hvp_method']
+        structure = settings['structure']
+        h = settings['h']
+
+        # ------------------------ calculate grad and hessian ------------------------ #
+        if hvp_method == 'autograd':
+            grad = vars.get_grad(create_graph=True)
+            def Hvp_fn1(x):
+                return hvp(params, grad, x, retain_graph=True)
+            Hvp_fn = Hvp_fn1
+
+        elif hvp_method == 'forward':
+            grad = vars.get_grad()
+            def Hvp_fn2(x):
+                return hvp_fd_forward(closure, params, x, h=h, g_0=grad, normalize=True)[1]
+            Hvp_fn = Hvp_fn2
+
+        elif hvp_method == 'central':
+            grad = vars.get_grad()
+            def Hvp_fn3(x):
+                return hvp_fd_central(closure, params, x, h=h, normalize=True)[1]
+            Hvp_fn = Hvp_fn3
+
+        else: raise ValueError(hvp_method)
+
+        # -------------------------------- inner step -------------------------------- #
+        update = vars.get_update()
+        if 'inner' in self.children:
+            update = apply(self.children['inner'], update, params=params, grads=grad, vars=vars)
+
+        # hessian
+        if structure.startswith('diagonal'):
+            H = Hvp_fn([torch.ones_like(p) for p in params])
+            if structure == 'diagonal1': torch._foreach_clamp_min_(H, 1)
+            if structure == 'diagonal_abs': torch._foreach_abs_(H)
+            torch._foreach_add_(H, reg)
+            torch._foreach_div_(update, H)
+            vars.update = update
+            return vars
+
+        # hessian
+        raise NotImplementedError(structure)
+
+
+
+
+

torchzero/modules/experimental/subspace_preconditioners.py

@@ -38,16 +38,19 @@ def apply_subspace_preconditioner(
     return basis @ update_projected # d
 
 class RandomSubspacePreconditioning(Transform):
-    """full matrix rmsprop in random subspace"""
-    def __init__(self, k: int, beta: float | None = 0.99):
-        defaults = dict(k=k, beta=beta)
+    """full matrix rmsprop in random slowly changing subspace"""
+    def __init__(self, k: int, beta: float | None = 0.99, basis_beta: float | None = 0.99, inner: Chainable | None = None):
+        defaults = dict(k=k, beta=beta, basis_beta=basis_beta)
         super().__init__(defaults, uses_grad=False)
 
+        if inner is not None: self.set_child('inner', inner)
+
     def transform(self, tensors, params, grads, vars):
         settings = self.settings[params[0]]
         g = torch.cat([t.view(-1) for t in tensors])
         k = settings['k']
         beta = settings['beta']
+        basis_beta = settings['basis_beta']
 
         if 'basis' not in self.global_state:
             self.global_state['basis'] = torch.randn(g.numel(), k, device=g.device, dtype=g.dtype)
@@ -56,13 +59,19 @@ class RandomSubspacePreconditioning(Transform):
         basis = self.global_state['basis']
         accumulator = self.global_state['accumulator']
 
+        if basis_beta is not None:
+            basis.lerp_(torch.randn_like(basis), 1-basis_beta)
+
         update_subspace_preconditioner_(g, basis, accumulator, beta)
+
+        if 'inner' in self.children:
+            tensors = apply(self.children['inner'], tensors, params, grads, vars)
+            g = torch.cat([t.view(-1) for t in tensors])
+
         try:
             preconditioned = apply_subspace_preconditioner(g, basis, accumulator)
         except torch.linalg.LinAlgError:
-            denom = g.abs().sum()
-            if denom <= 1e-10: denom = torch.ones_like(denom)
-            preconditioned = g / g.abs().sum()
+            preconditioned = g.clip(-0.1, 0.1)
         vec_to_tensors_(preconditioned, tensors)
 
         return tensors
@@ -119,9 +128,7 @@ class HistorySubspacePreconditioning(Transform):
         try:
             preconditioned = apply_subspace_preconditioner(g, basis, accumulator)
         except torch.linalg.LinAlgError:
-            denom = g.abs().sum()
-            if denom <= 1e-10: denom = torch.ones_like(denom)
-            preconditioned = g / g.abs().sum()
+            preconditioned = g.clip(-0.1,0.1)
         vec_to_tensors_(preconditioned, tensors)
 
         return tensors

torchzero/modules/optimizers/soap.py

@@ -222,8 +222,7 @@ class SOAP(Transform):
                     state['Q'] = get_orthogonal_matrix(state['GG'])
 
                 state['step'] = 0
-                updates.append(tensors[i].sign().div_(10))
-                # updates.append(tensors[i] / tensors[i].abs().sum())
+                updates.append(tensors[i].clip(-0.1, 0.1))
                 continue  # skip 1st step as in https://github.com/nikhilvyas/SOAP/blob/main/soap.py ?
                 # I use scaled update instead as to not mess up with next modules.
 

torchzero/modules/projections/projection.py

@@ -1,4 +1,5 @@
 import math
+from functools import partial
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
 from typing import Any, Literal
@@ -33,6 +34,25 @@ def _make_projected_closure(closure, vars: Vars, projection: "Projection",
 
     return projected_closure
 
+def _projected_get_grad_override(
+    retain_graph: bool | None = None,
+    create_graph: bool = False,
+    projection: Any = ...,
+    unprojected_vars: Any = ...,
+    self: Any = ...,
+):
+    assert isinstance(projection, Projection)
+    assert isinstance(unprojected_vars, Vars)
+    assert isinstance(self, Vars)
+
+    if self.grad is not None: return self.grad
+    grads = unprojected_vars.get_grad(retain_graph, create_graph)
+    projected_grads = list(projection.project(grads, self, current='grads'))
+    self.grad = projected_grads
+    for p, g in zip(self.params, projected_grads):
+        p.grad = g
+    return self.grad
+
 
 class Projection(Module, ABC):
     """
@@ -137,6 +157,12 @@ class Projection(Module, ABC):
 
         # step
         projected_vars.params = self._projected_params
+        projected_vars.get_grad = partial(
+            _projected_get_grad_override,
+            projection=self,
+            unprojected_vars=vars,
+            self=projected_vars,
+        )
         projected_vars = self.children['modules'].step(projected_vars)
 
         # empty fake params storage
@@ -149,7 +175,7 @@ class Projection(Module, ABC):
         unprojected_vars = projected_vars.clone(clone_update=False)
         unprojected_vars.closure = vars.closure
         unprojected_vars.params = vars.params
-        if unprojected_vars.grad is None: unprojected_vars.grad = vars.grad
+        unprojected_vars.grad = vars.grad
 
         if self._project_update:
             assert projected_vars.update is not None

torchzero/modules/quasi_newton/cg.py

@@ -64,7 +64,7 @@ class ConguateGradientBase(Transform, ABC):
 # ------------------------------- Polak-Ribière ------------------------------ #
 def polak_ribiere_beta(g: TensorList, prev_g: TensorList):
     denom = prev_g.dot(prev_g)
-    if denom == 0: return 0
+    if denom.abs() <= torch.finfo(g[0].dtype).eps: return 0
     return g.dot(g - prev_g) / denom
 
 class PolakRibiere(ConguateGradientBase):
@@ -76,8 +76,8 @@ class PolakRibiere(ConguateGradientBase):
         return polak_ribiere_beta(g, prev_g)
 
 # ------------------------------ Fletcher–Reeves ----------------------------- #
-def fletcher_reeves_beta(gg, prev_gg):
-    if prev_gg == 0: return 0
+def fletcher_reeves_beta(gg: torch.Tensor, prev_gg: torch.Tensor):
+    if prev_gg.abs() <= torch.finfo(gg.dtype).eps: return 0
     return gg / prev_gg
 
 class FletcherReeves(ConguateGradientBase):
@@ -98,7 +98,7 @@ class FletcherReeves(ConguateGradientBase):
 def hestenes_stiefel_beta(g: TensorList, prev_d: TensorList,prev_g: TensorList):
     grad_diff = g - prev_g
     denom = prev_d.dot(grad_diff)
-    if denom == 0: return 0
+    if denom.abs() < torch.finfo(g[0].dtype).eps: return 0
     return (g.dot(grad_diff) / denom).neg()
 
 
@@ -114,7 +114,7 @@ class HestenesStiefel(ConguateGradientBase):
 # --------------------------------- Dai–Yuan --------------------------------- #
 def dai_yuan_beta(g: TensorList, prev_d: TensorList,prev_g: TensorList):
     denom = prev_d.dot(g - prev_g)
-    if denom == 0: return 0
+    if denom.abs() <= torch.finfo(g[0].dtype).eps: return 0
     return (g.dot(g) / denom).neg()
 
 class DaiYuan(ConguateGradientBase):
@@ -129,7 +129,7 @@ class DaiYuan(ConguateGradientBase):
 # -------------------------------- Liu-Storey -------------------------------- #
 def liu_storey_beta(g:TensorList, prev_d:TensorList, prev_g:TensorList, ):
     denom = prev_g.dot(prev_d)
-    if denom == 0: return 0
+    if denom.abs() <= torch.finfo(g[0].dtype).eps: return 0
     return g.dot(g - prev_g) / denom
 
 class LiuStorey(ConguateGradientBase):
@@ -159,7 +159,7 @@ class ConjugateDescent(Transform):
             self.global_state['denom'] = torch.tensor(0.).to(g[0])
 
         prev_gd = self.global_state.get('prev_gd', 0)
-        if prev_gd == 0: beta = 0
+        if abs(prev_gd) <= torch.finfo(g[0].dtype).eps: beta = 0
         else: beta = g.dot(g) / prev_gd
 
         # inner step
@@ -176,7 +176,7 @@ class ConjugateDescent(Transform):
 def hager_zhang_beta(g:TensorList, prev_d:TensorList, prev_g:TensorList,):
     g_diff = g - prev_g
     denom = prev_d.dot(g_diff)
-    if denom == 0: return 0
+    if denom.abs() <= torch.finfo(g[0].dtype).eps: return 0
 
     term1 = 1/denom
     # term2
@@ -198,7 +198,7 @@ class HagerZhang(ConguateGradientBase):
 def hs_dy_beta(g: TensorList, prev_d: TensorList,prev_g: TensorList):
     grad_diff = g - prev_g
     denom = prev_d.dot(grad_diff)
-    if denom == 0: return 0
+    if denom.abs() <= torch.finfo(g[0].dtype).eps: return 0
 
     # Dai-Yuan
     dy_beta = (g.dot(g) / denom).neg().clamp(min=0)

torchzero/modules/quasi_newton/experimental/modular_lbfgs.py

@@ -37,10 +37,11 @@ def lbfgs(
     z_tfm: Any,
 ):
     if len(s_history) == 0 or y_k is None or ys_k is None:
-        # dir = params.grad.sign() # may work fine
 
-        # initial step size guess taken from pytorch L-BFGS
-        return tensors_.mul_(min(1.0, 1.0 / tensors_.abs().global_sum())) # pyright: ignore[reportArgumentType]
+        # initial step size guess modified from pytorch L-BFGS
+        scale = 1 / tensors_.abs().global_sum()
+        if scale < 1e-5: scale = 1 / tensors_.abs().mean()
+        return tensors_.mul_(min(1.0, scale)) # pyright: ignore[reportArgumentType]
 
     else:
         # 1st loop

torchzero/modules/quasi_newton/lbfgs.py

@@ -36,10 +36,11 @@ def lbfgs(
     step: int,
 ):
     if len(s_history) == 0 or y_k is None or ys_k is None:
-        # dir = params.grad.sign() # may work fine
 
-        # initial step size guess taken from pytorch L-BFGS
-        return tensors_.mul_(min(1.0, 1.0 / tensors_.abs().global_sum())) # pyright: ignore[reportArgumentType]
+        # initial step size guess modified from pytorch L-BFGS
+        scale_factor = 1 / TensorList(tensors_).abs().global_sum().clip(min=1)
+        scale_factor = scale_factor.clip(min=torch.finfo(tensors_[0].dtype).eps)
+        return tensors_.mul_(scale_factor)
 
     else:
         # 1st loop

torchzero/modules/quasi_newton/lsr1.py

@@ -17,8 +17,9 @@ def lsr1_(
 ):
     if step == 0 or not s_history:
         # initial step size guess from pytorch
-        tensors_.div_(max(1.0, tensors_.abs().global_sum())) # pyright:ignore[reportArgumentType]
-        return tensors_
+        scale_factor = 1 / TensorList(tensors_).abs().global_sum().clip(min=1)
+        scale_factor = scale_factor.clip(min=torch.finfo(tensors_[0].dtype).eps)
+        return tensors_.mul_(scale_factor)
 
     m = len(s_history)
 
@@ -64,7 +65,10 @@ def lsr1_(
         Hx.add_(w_k, alpha=w_k.dot(tensors_) / wy) # pyright:ignore[reportArgumentType]
 
     if scale_second and step == 1:
-        Hx.div_(max(1.0, tensors_.abs().global_sum())) # pyright:ignore[reportArgumentType]
+        scale_factor = 1 / TensorList(tensors_).abs().global_sum().clip(min=1)
+        scale_factor = scale_factor.clip(min=torch.finfo(tensors_[0].dtype).eps)
+        Hx.mul_(scale_factor)
+
     return Hx
 
 

torchzero/modules/quasi_newton/quasi_newton.py

@@ -68,6 +68,7 @@ class HessianUpdateStrategy(TensorwisePreconditioner, ABC):
         M_key = 'H' if inverse else 'B'
         M = state.get(M_key, None)
         step = state.get('step', 0)
+        state['step'] = step + 1
         init_scale = settings['init_scale']
         tol = settings['tol']
         tol_reset = settings['tol_reset']
@@ -91,13 +92,12 @@ class HessianUpdateStrategy(TensorwisePreconditioner, ABC):
         state['p_prev'].copy_(p)
         state['g_prev'].copy_(g)
 
-
-        if reset_interval is not None and step % reset_interval == 0:
+        if reset_interval is not None and step != 0 and step % reset_interval == 0:
             self._reset_M_(M, s, y, inverse, init_scale)
             return
 
         # tolerance on gradient difference to avoid exploding after converging
-        if y.abs().max() <= tol:
+        elif y.abs().max() <= tol:
             # reset history
             if tol_reset: self._reset_M_(M, s, y, inverse, init_scale)
             return
@@ -119,11 +119,12 @@ class HessianUpdateStrategy(TensorwisePreconditioner, ABC):
 
     @torch.no_grad
     def apply_tensor(self, tensor, param, grad, state, settings):
-        step = state['step'] = state.get('step', 0) + 1
+        step = state.get('step', 0)
 
         if settings['scale_second'] and step == 2:
-            s = max(1, tensor.abs().sum()) # pyright:ignore[reportArgumentType]
-            if s < settings['tol']: tensor = tensor/s
+            scale_factor = 1 / tensor.abs().sum().clip(min=1)
+            scale_factor = scale_factor.clip(min=torch.finfo(tensor.dtype).eps)
+            tensor = tensor * scale_factor
 
         inverse = settings['inverse']
         if inverse:
@@ -135,7 +136,7 @@ class HessianUpdateStrategy(TensorwisePreconditioner, ABC):
         return torch.linalg.solve_ex(B, tensor.view(-1))[0].view_as(tensor) # pylint:disable=not-callable
 
 # to avoid typing all arguments for each method
-class QuasiNewtonH(HessianUpdateStrategy):
+class HUpdateStrategy(HessianUpdateStrategy):
     def __init__(
         self,
         init_scale: float | Literal["auto"] = "auto",
@@ -174,7 +175,7 @@ def bfgs_H_(H:torch.Tensor, s: torch.Tensor, y:torch.Tensor, tol: float):
     H += term1.sub_(term2)
     return H
 
-class BFGS(QuasiNewtonH):
+class BFGS(HUpdateStrategy):
     def update_H(self, H, s, y, p, g, p_prev, g_prev, state, settings):
         return bfgs_H_(H=H, s=s, y=y, tol=settings['tol'])
 
@@ -193,7 +194,7 @@ def sr1_H_(H:torch.Tensor, s: torch.Tensor, y:torch.Tensor, tol:float):
     H += torch.outer(z, z).div_(denom)
     return H
 
-class SR1(QuasiNewtonH):
+class SR1(HUpdateStrategy):
     def update_H(self, H, s, y, p, g, p_prev, g_prev, state, settings):
         return sr1_H_(H=H, s=s, y=y, tol=settings['tol'])
 
@@ -213,7 +214,7 @@ def dfp_H_(H:torch.Tensor, s: torch.Tensor, y:torch.Tensor, tol: float):
     H += term1.sub_(term2)
     return H
 
-class DFP(QuasiNewtonH):
+class DFP(HUpdateStrategy):
     def update_H(self, H, s, y, p, g, p_prev, g_prev, state, settings):
         return dfp_H_(H=H, s=s, y=y, tol=settings['tol'])
 
@@ -254,19 +255,19 @@ def greenstadt2_H_(H:torch.Tensor, s: torch.Tensor, y:torch.Tensor, tol: float):
     H -= num/denom
     return H
 
-class BroydenGood(QuasiNewtonH):
+class BroydenGood(HUpdateStrategy):
     def update_H(self, H, s, y, p, g, p_prev, g_prev, state, settings):
         return broyden_good_H_(H=H, s=s, y=y, tol=settings['tol'])
 
-class BroydenBad(QuasiNewtonH):
+class BroydenBad(HUpdateStrategy):
     def update_H(self, H, s, y, p, g, p_prev, g_prev, state, settings):
         return broyden_bad_H_(H=H, s=s, y=y, tol=settings['tol'])
 
-class Greenstadt1(QuasiNewtonH):
+class Greenstadt1(HUpdateStrategy):
     def update_H(self, H, s, y, p, g, p_prev, g_prev, state, settings):
         return greenstadt1_H_(H=H, s=s, y=y, g_prev=g_prev, tol=settings['tol'])
 
-class Greenstadt2(QuasiNewtonH):
+class Greenstadt2(HUpdateStrategy):
     def update_H(self, H, s, y, p, g, p_prev, g_prev, state, settings):
         return greenstadt2_H_(H=H, s=s, y=y, tol=settings['tol'])
 
@@ -287,7 +288,7 @@ def column_updating_H_(H:torch.Tensor, s:torch.Tensor, y:torch.Tensor, tol:float
     H[:, j] += num.squeeze() / denom
     return H
 
-class ColumnUpdatingMethod(QuasiNewtonH):
+class ColumnUpdatingMethod(HUpdateStrategy):
     """Lopes, V. L., & Martínez, J. M. (1995). Convergence properties of the inverse column-updating method. Optimization Methods & Software, 6(2), 127–144. from https://www.ime.unicamp.br/sites/default/files/pesquisa/relatorios/rp-1993-76.pdf"""
     def update_H(self, H, s, y, p, g, p_prev, g_prev, state, settings):
         return column_updating_H_(H=H, s=s, y=y, tol=settings['tol'])
@@ -307,7 +308,7 @@ def thomas_H_(H: torch.Tensor, R:torch.Tensor, s: torch.Tensor, y: torch.Tensor,
     H -= num/denom
     return H, R
 
-class ThomasOptimalMethod(QuasiNewtonH):
+class ThomasOptimalMethod(HUpdateStrategy):
     """Thomas, Stephen Walter. Sequential estimation techniques for quasi-Newton algorithms. Cornell University, 1975."""
     def update_H(self, H, s, y, p, g, p_prev, g_prev, state, settings):
         if 'R' not in state: state['R'] = torch.eye(H.size(-1), device=H.device, dtype=H.dtype)
@@ -364,7 +365,7 @@ def pearson2_H_(H:torch.Tensor, s: torch.Tensor, y:torch.Tensor, tol: float):
     H += num.div_(sy)
     return H
 
-class Pearson2(QuasiNewtonH):
+class Pearson2(HUpdateStrategy):
     """finally found a reference in https://www.recotechnologies.com/~beigi/ps/asme-jdsmc-93-2.pdf"""
     def update_H(self, H, s, y, p, g, p_prev, g_prev, state, settings):
         return pearson2_H_(H=H, s=s, y=y, tol=settings['tol'])

torchzero/modules/second_order/__init__.py

@@ -1,3 +1,3 @@
 from .newton import Newton
 from .newton_cg import NewtonCG
-from .nystrom import NystromSketchAndSolve, NystromPCG
+from .nystrom import NystromSketchAndSolve, NystromPCG

torchzero/modules/second_order/newton.py

@@ -1,14 +1,18 @@
 import warnings
+from collections.abc import Callable
 from functools import partial
 from typing import Literal
-from collections.abc import Callable
+
 import torch
 
-from ...core import Chainable, apply, Module
-from ...utils import vec_to_tensors, TensorList
+from ...core import Chainable, Module, apply
+from ...utils import TensorList, vec_to_tensors
 from ...utils.derivatives import (
     hessian_list_to_mat,
     hessian_mat,
+    hvp,
+    hvp_fd_central,
+    hvp_fd_forward,
     jacobian_and_hessian_wrt,
 )
 
@@ -117,9 +121,10 @@ class Newton(Module):
             raise ValueError(hessian_method)
 
         # -------------------------------- inner step -------------------------------- #
+        update = vars.get_update()
         if 'inner' in self.children:
-            g_list = apply(self.children['inner'], list(g_list), params=params, grads=list(g_list), vars=vars)
-        g = torch.cat([t.view(-1) for t in g_list])
+            update = apply(self.children['inner'], update, params=params, grads=list(g_list), vars=vars)
+        g = torch.cat([t.view(-1) for t in update])
 
         # ------------------------------- regulazition ------------------------------- #
         if eig_reg: H = eig_tikhonov_(H, reg)
@@ -139,4 +144,4 @@ class Newton(Module):
         if update is None: update = least_squares_solve(H, g)
 
         vars.update = vec_to_tensors(update, params)
-        return vars
+        return vars

torchzero/modules/second_order/newton_cg.py

@@ -66,9 +66,9 @@ class NewtonCG(Module):
 
 
         # -------------------------------- inner step -------------------------------- #
-        b = grad
+        b = vars.get_update()
         if 'inner' in self.children:
-            b = as_tensorlist(apply(self.children['inner'], [g.clone() for g in grad], params=params, grads=grad, vars=vars))
+            b = as_tensorlist(apply(self.children['inner'], b, params=params, grads=grad, vars=vars))
 
         # ---------------------------------- run cg ---------------------------------- #
         x0 = None
@@ -76,7 +76,7 @@ class NewtonCG(Module):
         x = cg(A_mm=H_mm, b=as_tensorlist(b), x0_=x0, tol=tol, maxiter=maxiter, reg=reg)
         if warm_start:
             assert x0 is not None
-            x0.set_(x)
+            x0.copy_(x)
 
         vars.update = x
         return vars

torchzero/modules/second_order/nystrom.py

@@ -15,7 +15,7 @@ class NystromSketchAndSolve(Module):
         rank: int,
         reg: float = 1e-3,
         hvp_method: Literal["forward", "central", "autograd"] = "autograd",
-        h=1e-3,
+        h=1e-2,
         inner: Chainable | None = None,
         seed: int | None = None,
     ):
@@ -74,9 +74,9 @@ class NystromSketchAndSolve(Module):
 
 
         # -------------------------------- inner step -------------------------------- #
-        b = grad
+        b = vars.get_update()
         if 'inner' in self.children:
-            b = apply(self.children['inner'], [g.clone() for g in grad], params=params, grads=grad, vars=vars)
+            b = apply(self.children['inner'], b, params=params, grads=grad, vars=vars)
 
         # ------------------------------ sketch&n&solve ------------------------------ #
         x = nystrom_sketch_and_solve(A_mm=H_mm, b=torch.cat([t.ravel() for t in b]), rank=rank, reg=reg, generator=generator)
@@ -93,7 +93,7 @@ class NystromPCG(Module):
         tol=1e-3,
         reg: float = 1e-6,
         hvp_method: Literal["forward", "central", "autograd"] = "autograd",
-        h=1e-3,
+        h=1e-2,
         inner: Chainable | None = None,
         seed: int | None = None,
     ):
@@ -156,9 +156,9 @@ class NystromPCG(Module):
 
 
         # -------------------------------- inner step -------------------------------- #
-        b = grad
+        b = vars.get_update()
         if 'inner' in self.children:
-            b = apply(self.children['inner'], [g.clone() for g in grad], params=params, grads=grad, vars=vars)
+            b = apply(self.children['inner'], b, params=params, grads=grad, vars=vars)
 
         # ------------------------------ sketch&n&solve ------------------------------ #
         x = nystrom_pcg(A_mm=H_mm, b=torch.cat([t.ravel() for t in b]), sketch_size=sketch_size, reg=reg, tol=tol, maxiter=maxiter, x0_=None, generator=generator)

torchzero/utils/linalg/benchmark.py

@@ -0,0 +1,20 @@
+from collections.abc import Callable
+
+import torch
+
+
+def benchmark_solver(
+    A: torch.Tensor | Callable[[torch.Tensor], torch.Tensor],
+    b: torch.Tensor,
+    solver: Callable[[Callable[[torch.Tensor], torch.Tensor], torch.Tensor]]
+):
+    residuals = []
+    def A_mm(x):
+        if callable(A): Ax = A(x)
+        else: Ax = A@x
+        residuals.append(torch.linalg.vector_norm(Ax-b)) # pylint:disable=not-callable
+        return Ax
+
+    solver(A_mm, b)
+    return residuals
+

torchzero/utils/linalg/solve.py

@@ -8,27 +8,27 @@ from .. import TensorList, generic_zeros_like, generic_vector_norm, generic_nume
 def cg(
     A_mm: Callable[[torch.Tensor], torch.Tensor],
     b: torch.Tensor,
-    x0_: torch.Tensor | None,
-    tol: float | None,
-    maxiter: int | None,
+    x0_: torch.Tensor | None = None,
+    tol: float | None = 1e-4,
+    maxiter: int | None = 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,
+    x0_: TensorList | None = None,
+    tol: float | None = 1e-4,
+    maxiter: int | None = 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,
+    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
@@ -90,7 +90,7 @@ def nystrom_sketch_and_solve(
     A_mm: Callable[[torch.Tensor], torch.Tensor],
     b: torch.Tensor,
     rank: int,
-    reg: float,
+    reg: float = 1e-3,
     generator=None,
 ) -> torch.Tensor:
     U, lambd = nystrom_approximation(
@@ -116,10 +116,10 @@ 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,
+    reg: float = 1e-6,
+    x0_: torch.Tensor | None = None,
+    tol: float | None = 1e-4,
+    maxiter: int | None = None,
     generator=None,
 ) -> torch.Tensor:
     U, lambd = nystrom_approximation(
@@ -166,3 +166,4 @@ def nystrom_pcg(
         z = P_inv @ residual
         beta = residual.dot(z) / rz
         p = z + p*beta
+

{torchzero-0.3.6.dist-info → torchzero-0.3.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchzero
-Version: 0.3.6
+Version: 0.3.9
 Summary: Modular optimization library for PyTorch.
 Author-email: Ivan Nikishev <nkshv2@gmail.com>
 License: MIT License
@@ -156,7 +156,7 @@ for epoch in range(100):
   * `Newton`: Classic Newton's method.
   * `NewtonCG`: Matrix-free newton's method with conjugate gradient solver.
   * `NystromSketchAndSolve`: Nyström sketch-and-solve method.
-  * `NystromPCG`: NewtonCG with Nyström preconditioning (my current recommendation).
+  * `NystromPCG`: NewtonCG with Nyström preconditioning (usually beats NewtonCG).
 
 * **Quasi-Newton**: Approximate second-order optimization methods.
   * `LBFGS`: Limited-memory BFGS.

{torchzero-0.3.6.dist-info → torchzero-0.3.9.dist-info}/RECORD

@@ -1,14 +1,14 @@
 docs/source/conf.py,sha256=jd80ZT2IdCx7nlQrpOTJL8UhGBNm6KYyXlpp0jmRiAw,1849
 tests/test_identical.py,sha256=NZ7A8Rm1U9Q16d-cG2G_wccpPtNALyoKYJt9qMownMc,11568
 tests/test_module.py,sha256=qX3rjdSJsbA8JO17bPTUIDspe7bg2dogqxMw__KV7SU,2039
-tests/test_opts.py,sha256=oDZVFr9AE9ZhyR-sImSgNzQsbPsUtJLzuLd1Nxgkp1w,40850
-tests/test_tensorlist.py,sha256=VWX9wYdfkG-0Y8I0wWPp56ZJM0mBNPvS_SC3irmcYcs,72427
+tests/test_opts.py,sha256=TZVaCv2ZLdHSkL6snTEkqhTMHqlcO55L-c56k6Hh4xc,40850
+tests/test_tensorlist.py,sha256=Djpr5C0T5d_gz-j-P-bpo_X51DC4twbtT9c-xDSFbP0,72438
 tests/test_utils_optimizer.py,sha256=bvC0Ehvs2L8fohpyIF5Vfr9OKTycpnODWLPflXilU1c,8414
 tests/test_vars.py,sha256=3p9dsHk7SJpMd-WRD0ziBNq5FEHRBJGSxbMLD8ES4J0,6815
 torchzero/__init__.py,sha256=L7IJ1qZ3o8E9oRwlJZBK2_2yII_eeGEk57Of6EfVbrk,112
 torchzero/core/__init__.py,sha256=2JRyeGZprTexAeEPQOIl9fLFGBwzvya-AwKyt7XAmGQ,210
 torchzero/core/module.py,sha256=Razw3c71Kfegznm0vQxsii1KuTUCPBC9UGyq2v-KX4M,27568
-torchzero/core/preconditioner.py,sha256=rMYusKbaypm5K0Ii9VdjKhxi2YWNQbBk9f6AV_MJulY,6191
+torchzero/core/preconditioner.py,sha256=n9oh7kZdt1kU3Wh472lnvLrsXwhR5Wqe6lIp7JuAJ_I,6336
 torchzero/core/transform.py,sha256=ajNJcX45ds-_lc5CqxgLfEFGil6_BYLerB0WvoTi8rM,10303
 torchzero/modules/__init__.py,sha256=BDeyuSd2s1WFUUXIo3tGTNp4aYp4A2B94cydpPW24nY,332
 torchzero/modules/functional.py,sha256=HXNzmPe7LsPadryEm7zrcEKqGej16QDwSgBkbEvggFM,6492
@@ -16,18 +16,20 @@ torchzero/modules/clipping/__init__.py,sha256=ZaffMF7mIRK6hZSfuZadgjNTX6hF5ANiLB
 torchzero/modules/clipping/clipping.py,sha256=I-5utyrqdKtF5yaH-9m2F3UqdfpPmA2bSSFUAZ_d60Q,12544
 torchzero/modules/clipping/ema_clipping.py,sha256=pLeNuEBLpJ74io2sHn_ZVYaQ6ydEfhpVfVEX2bFttd0,5947
 torchzero/modules/clipping/growth_clipping.py,sha256=OD-kdia2Rn-DvYlYV6EZlGPDVTh9tj-W9mpiZPc3hOQ,6772
-torchzero/modules/experimental/__init__.py,sha256=sJ6URgX35P3zJ2ugBKgAcwBWmdBmAPDW3vXHQ0sK-ro,443
-torchzero/modules/experimental/absoap.py,sha256=XUHr5SeLdhLW2kMvWea5xAqZeuJBDQoO4zprDxs4bgU,13317
-torchzero/modules/experimental/adadam.py,sha256=W7rRXYJ9tGrzqD_FdFX00HBLuWOEr2tHtfshf6lDFYE,4049
-torchzero/modules/experimental/adamY.py,sha256=FoSn-qMI5_BdqZH10WGKkl-zYTPESBdGZ9lfhyqnbB0,4591
-torchzero/modules/experimental/adasoap.py,sha256=07gPdEdBIKtmdmSzTGtTO0c2ZkS_otVLufQ76okBjHY,11239
-torchzero/modules/experimental/algebraic_newton.py,sha256=_XFYR6bdHWgA5bozxc9AJYteBIAnHrSLgo_bSaZ13eg,5193
-torchzero/modules/experimental/curveball.py,sha256=Nw9jtSp5QNj7-FN3qshjYEDHc68LwRLha-Co78mfR5w,3242
-torchzero/modules/experimental/dsoap.py,sha256=BEZDw3_n5VDhu7VLgkoSN4rI9JeBdGoO9gFZfqsh74M,10983
-torchzero/modules/experimental/gradmin.py,sha256=55dpBDNyrYJusluFhw-v1BXuj1UxER7pNEPTtwYKD4E,3648
+torchzero/modules/experimental/__init__.py,sha256=fEPDYDl7qhaFoferDRmG3ehwuqSvx4Vt2uOz0Y7h4to,483
+torchzero/modules/experimental/absoap.py,sha256=Z4MS4pDPSQ9IaTk8g57OfrsWcYVOT72x533KKtn2Zxk,13512
+torchzero/modules/experimental/adadam.py,sha256=OAPF1-NUbg79V3QOTYzsQlRC97C7XHj5boOLDqLz3PE,4029
+torchzero/modules/experimental/adamY.py,sha256=g1pAHwgdyDdKvObZ67lCSc36L99tl5jlQgOr4lMJCDo,4595
+torchzero/modules/experimental/adasoap.py,sha256=JdV6rB9xfqL3vbHpZCLmkJZKRObZ1nVoEmabtIeVT3E,11195
+torchzero/modules/experimental/algebraic_newton.py,sha256=sq5ZD_j_EtlxIjNnS0rKKwTSG_JuwsZOg9ZMMQTuQm0,5154
+torchzero/modules/experimental/curveball.py,sha256=Uk30uLEztTHD5IUJLJm9Nn3x31DF9kQHmeLFhc065us,3262
+torchzero/modules/experimental/gradmin.py,sha256=iJmEvDEdVdck0C-94pY3iGxnIoNv6Fu6vj3f7lS6aQM,3686
+torchzero/modules/experimental/newton_solver.py,sha256=iGI2LHLaZd2ovpbq1Vogs76os0zWG7VwM7nUz8RzxVg,3071
 torchzero/modules/experimental/reduce_outward_lr.py,sha256=kjtRwepBGBca77ToM-lw3b8ywptMtmSdC_jQfjJAwlY,1184
-torchzero/modules/experimental/spectral.py,sha256=D3_nCI8teFirCdnnLprNnZ3G1gsOB6RUBWCeDbwi7P0,12043
-torchzero/modules/experimental/subspace_preconditioners.py,sha256=4SRJOyTG-fJCGunHR62aRrzw3qFmeI6fRQAYHIadhWw,4682
+torchzero/modules/experimental/soapy.py,sha256=Ishd2Jj6BbhjrLyC48zf-cjMmA1kJb_uKXESQBIML_s,10990
+torchzero/modules/experimental/spectral.py,sha256=8_n208V2yPY3z5pCym-FvwO7DGFhozNgWlpIBtQSdrI,12139
+torchzero/modules/experimental/structured_newton.py,sha256=uWczR-uAXHaFwf0mlOThv2sLG0irH6Gz1hKlGHtPAj4,3386
+torchzero/modules/experimental/subspace_preconditioners.py,sha256=WnHpga7Kx4-N2xU5vP3uUHRER70ymyNJCWbSx2zXWOk,4976
 torchzero/modules/experimental/tropical_newton.py,sha256=uq66ouhgrgc8iYGozDQ3_rtbubj8rKRwb1jfcdnlpHg,4903
 torchzero/modules/grad_approximation/__init__.py,sha256=DVFjf0cXuF70NA0nJ2WklpP01PQgrRZxUjUQjjQeSos,195
 torchzero/modules/grad_approximation/fdm.py,sha256=2PNNBIMup1xlOwLFAwAS3xAVd-7GGVyerMeKH1ug9LQ,3591
@@ -70,26 +72,26 @@ torchzero/modules/optimizers/orthograd.py,sha256=5BLnNJTYuGUClHmlxaXZ1jNvBR4zSFD
 torchzero/modules/optimizers/rmsprop.py,sha256=d10Y9Ck-391tVysO3xMHg3g2Pe0UEZplgebEyDYi3Z4,4333
 torchzero/modules/optimizers/rprop.py,sha256=n4k5-9F3ppH0Xl-4l4vNXfqVf2r67vMPCkstUaQKPLw,10974
 torchzero/modules/optimizers/shampoo.py,sha256=AHHV6d71DqKDPCg52ShWIPIRSGtWkMc1v1XwXgDG3qY,8606
-torchzero/modules/optimizers/soap.py,sha256=HL1YrfiEiRMh6aW9D5UEZXBjo3yMTqnpKPHXVD8fOa8,11590
+torchzero/modules/optimizers/soap.py,sha256=Kf2BAtIf2QY1V2ZJcUjRLcp2WfIVLd3mNclnaT3Nmds,11520
 torchzero/modules/optimizers/sophia_h.py,sha256=8pSlYVm66xWplzdP8MX3MCTzzIYHsxGzDEXJKA03Zgg,4279
 torchzero/modules/projections/__init__.py,sha256=OCxlh_-Tx-xpl31X03CeFJH9XveH563oEsWc8rUvX0A,196
 torchzero/modules/projections/dct.py,sha256=wxaEV6dTNiOqW_n2UHX0De6mMXTKDXK6UNcMNI4Rogk,2373
 torchzero/modules/projections/fft.py,sha256=OpCcEM1-A2dgk1umwRsBsvK7ObiHtsBKlkkcw0IX83Q,2961
 torchzero/modules/projections/galore.py,sha256=c9CZ0kHxpKEoyfc_lnmeHOkNp55jCppb7onN5YmWnN8,242
-torchzero/modules/projections/projection.py,sha256=tvUBZ4XGY1GkOg6jrKS7FvpIpjUc2FJL_SMRpoROT1E,9330
+torchzero/modules/projections/projection.py,sha256=aYufSD3ftRUqVScPmqxwEFgP1P8ioxM8z9eyzaL7d10,10147
 torchzero/modules/projections/structural.py,sha256=QaCGHmzHCXj46sM-XZ5XlYU9BnuRKI2ReR3LE8y2R4g,5740
 torchzero/modules/quasi_newton/__init__.py,sha256=0iOlX73PHj9lQS3_2cJ5lyCdas904MnFfIvR8Popvzw,402
-torchzero/modules/quasi_newton/cg.py,sha256=h-di1oKKP1tDoh-LogBRIRCp2UF9GA6XjEJPlX6xXf4,9322
-torchzero/modules/quasi_newton/lbfgs.py,sha256=jtO5ldbx66yUWv-20c-4mvq6HhCMuomCwJK8A8bjcYA,9168
-torchzero/modules/quasi_newton/lsr1.py,sha256=F_DtMQZfQSjmSLjnx4nw16AV7qCdNxT9ITQbfNFrPdM,5879
+torchzero/modules/quasi_newton/cg.py,sha256=lIJvfWAZ08r0o4uqaJnRG6pvcE2kBkJUkZ1MK37KMTk,9602
+torchzero/modules/quasi_newton/lbfgs.py,sha256=SMgesPMZ4ubVeG7R395SnAb5ffkyPHbzSQMqPlLGI7U,9211
+torchzero/modules/quasi_newton/lsr1.py,sha256=XmYyYANzQgQuFtOMW59znQrS-mprGRXazicfB9JAup8,6059
 torchzero/modules/quasi_newton/olbfgs.py,sha256=2YAOXlMnPGw22sNcIMH1hmggzAXQRbN59RSPUZNKUZY,8352
-torchzero/modules/quasi_newton/quasi_newton.py,sha256=jwQkzlnozIaxHW9kuDAAlME0YuQdrdZX9OZZoTmej4Q,17384
+torchzero/modules/quasi_newton/quasi_newton.py,sha256=rUp4s3MbACcOjwpz00TAjl-olif50voTmC16vv5XrSE,17496
 torchzero/modules/quasi_newton/experimental/__init__.py,sha256=3qpZGgdsx6wpoafWaNWx-eamRl1FuxVCWQZq8Y7Cl98,39
-torchzero/modules/quasi_newton/experimental/modular_lbfgs.py,sha256=PlyuIH2pFazIR89OGTrZESt752GkbArh_Zb8mtVCOi0,10731
-torchzero/modules/second_order/__init__.py,sha256=5lRmwIU53eRc1owpOZ5FMDc7u1Z48I3PDc0NyCBaJNM,113
-torchzero/modules/second_order/newton.py,sha256=XNhscAuWwxOUwps3sUrxc2ExgkNFbilnAdszrCvQxFg,5845
-torchzero/modules/second_order/newton_cg.py,sha256=lUVn4-ZoW3qAxqEy8i7yz_aN7sZDoQChd-A_Ubrz-Ag,2871
-torchzero/modules/second_order/nystrom.py,sha256=ZyCWrde-_-Ednj46jafuvBOzG3nC-3cPYGr-HytZbsE,6073
+torchzero/modules/quasi_newton/experimental/modular_lbfgs.py,sha256=ec6JKYX89xA_UlY9VrMB3hBjDyNKwkalS_4JQGA1qOY,10762
+torchzero/modules/second_order/__init__.py,sha256=jolCGaIVkID9hpxgx0Tc22wgjVlwuWekWjKTMe5jKXw,114
+torchzero/modules/second_order/newton.py,sha256=xxkrhFK4i5I9oOX3AGGh_6bXNDUSFq4D0pw3c7qgEd8,5925
+torchzero/modules/second_order/newton_cg.py,sha256=PILHRf2koop_cywE1RNGukT16alDO7prC4C3HlZcW30,2861
+torchzero/modules/second_order/nystrom.py,sha256=zdLSTQ_S5VViUt2sAmFNoDCCHKmHP2A7112czkZNlUk,6051
 torchzero/modules/smoothing/__init__.py,sha256=tUTGN0A-EQC7xuLV2AuHFWk-t7D6jIJlpV_3qyfRqLk,80
 torchzero/modules/smoothing/gaussian.py,sha256=YlT_G4MqAVkiWG56RHAwgt5SSPISpvQZQbSLh8mhF3I,6153
 torchzero/modules/smoothing/laplacian.py,sha256=Bfrs7D59SfdU7j-97UBKD1hs0obC-ZgjJvG7oKwaa0o,5065
@@ -116,13 +118,14 @@ torchzero/utils/python_tools.py,sha256=RFBqNj8w52dpJ983pUPPDbg2x1MX_-SsBnBMffWGG
 torchzero/utils/tensorlist.py,sha256=qSbiliVo1euFAksdHHHRbPUdYYxfkw1dvhpXj71wGy0,53162
 torchzero/utils/torch_tools.py,sha256=ohqnnZRlqdfp5PAfMSbQDIEKygW0_ARjxSEBp3Zo9nU,4756
 torchzero/utils/linalg/__init__.py,sha256=Dzbho3_z7JDdKzYD-QdLArg0ZEoC2BVGdlE3JoAnXHQ,272
+torchzero/utils/linalg/benchmark.py,sha256=wiIMn-GY2xxWbHVf8CPbJddUPeUPq9OUDkvbp1iILYI,479
 torchzero/utils/linalg/matrix_funcs.py,sha256=-LecWrPWbJvfeCgIzUhfWARa2aSZvJ12lHX7Jno38O4,3099
 torchzero/utils/linalg/orthogonalize.py,sha256=mDCkET7qgDZqf_y6oPYAK3d2L5HrB8gzOFPl0YoONaY,399
 torchzero/utils/linalg/qr.py,sha256=L-RXuYV-SIHI-Llq4y1rQ_Tz-yamds0_QNZeHapbjNE,2507
-torchzero/utils/linalg/solve.py,sha256=hN450ONzAirYOvWF2g0E0Wy2n1bCw4X-KXWi6p4jvDM,5136
+torchzero/utils/linalg/solve.py,sha256=P0PMi0zro3G3Rd0X-JeoLk7tqYDB0js0aB4bpQ0OABU,5235
 torchzero/utils/linalg/svd.py,sha256=wBxl-JSciINV-N6zvM4SGdveqMr6idq51h68LyQQRYg,660
-torchzero-0.3.6.dist-info/licenses/LICENSE,sha256=r9ZciAoZoqKC_FNADE0ORukj1p1XhLXEbegdsAyqhJs,1087
-torchzero-0.3.6.dist-info/METADATA,sha256=wjXJuO_WRQYv15BSA_9yo2qe2xe7jET7YOy8xb9YmnE,13944
-torchzero-0.3.6.dist-info/WHEEL,sha256=zaaOINJESkSfm_4HQVc5ssNzHCPXhJm0kEUakpsEHaU,91
-torchzero-0.3.6.dist-info/top_level.txt,sha256=YDdpIOb7HyKV9THOtOYsFFMTbxvCO0kiol4-83tDj-A,21
-torchzero-0.3.6.dist-info/RECORD,,
+torchzero-0.3.9.dist-info/licenses/LICENSE,sha256=r9ZciAoZoqKC_FNADE0ORukj1p1XhLXEbegdsAyqhJs,1087
+torchzero-0.3.9.dist-info/METADATA,sha256=aENIaMgy94tD6nakRWfApleVSy6bxW8-q3-mQeVSeGA,13941
+torchzero-0.3.9.dist-info/WHEEL,sha256=zaaOINJESkSfm_4HQVc5ssNzHCPXhJm0kEUakpsEHaU,91
+torchzero-0.3.9.dist-info/top_level.txt,sha256=YDdpIOb7HyKV9THOtOYsFFMTbxvCO0kiol4-83tDj-A,21
+torchzero-0.3.9.dist-info/RECORD,,