torchzero 0.3.6__tar.gz → 0.3.9__tar.gz

{torchzero-0.3.6 → torchzero-0.3.9}/PKG-INFO

@@ -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 → torchzero-0.3.9}/README.md

@@ -117,7 +117,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 → torchzero-0.3.9}/pyproject.toml

@@ -2,7 +2,7 @@
 # STEP 1 - COMMIT NEW CHANGES BUT DON'T PUSH THEM YET
 # STEP 2 - BUMP VERSION AND COMMIT IT (DONT PUSH!!!!)
 # STEP 3 - CREATE TAG WITH THAT VERSION
-# STEP 4 - PUSH CHANGES
+# STEP 4 - PUSH (SYNC) CHANGES
 # STEP 5 - PUSH TAG
 
 [build-system]
@@ -13,7 +13,7 @@ build-backend = "setuptools.build_meta"
 name = "torchzero"
 description = "Modular optimization library for PyTorch."
 
-version = "0.3.6"
+version = "0.3.9"
 dependencies = [
   "torch",
   "numpy",

{torchzero-0.3.6 → torchzero-0.3.9}/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,
 )
 

{torchzero-0.3.6 → torchzero-0.3.9}/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-0.3.6 → torchzero-0.3.9}/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-0.3.6 → torchzero-0.3.9}/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-0.3.6 → torchzero-0.3.9}/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-0.3.6 → torchzero-0.3.9}/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-0.3.6 → torchzero-0.3.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-0.3.6 → torchzero-0.3.9}/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-0.3.6 → torchzero-0.3.9}/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-0.3.6 → torchzero-0.3.9}/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-0.3.6 → torchzero-0.3.9}/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-0.3.9/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-0.3.6/torchzero/modules/experimental/dsoap.py → torchzero-0.3.9/torchzero/modules/experimental/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-0.3.6 → torchzero-0.3.9}/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-0.3.9/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-0.3.6 → torchzero-0.3.9}/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-0.3.6 → torchzero-0.3.9}/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-0.3.6 → torchzero-0.3.9}/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