torchzero 0.4.1__py3-none-any.whl → 0.4.2__py3-none-any.whl

torchzero/core/module.py

@@ -2,18 +2,19 @@ import warnings
 from abc import ABC, abstractmethod
 from collections import ChainMap, defaultdict
 from collections.abc import Callable, Iterable, Sequence
-from typing import Any, overload, TYPE_CHECKING
+from typing import Any, overload, TYPE_CHECKING, Literal
 
 import torch
 
 from ..linalg.linear_operator import LinearOperator
 from ..utils.optimizer import Init, ListLike, get_state_vals
-from ..utils.params import Params, _make_param_groups
+from ..utils.params import Params, _make_param_groups, _set_fake_params_, _empty_fake_param_storage_
 from .functional import step_tensors
 
 if TYPE_CHECKING:
     from .objective import Objective
 
+ProjectedBuffer = Literal["grad", "grad_sq", "grad_cu", "covariance", "inverse"]
 
 class Module(ABC):
     """Abstract base class for an optimizer modules.
@@ -52,6 +53,12 @@ class Module(ABC):
         self._overridden_keys = set()
         """tracks keys overridden with ``set_param_groups``, only used to not give a warning"""
 
+        self._projected_keys: defaultdict[ProjectedBuffer, set[str]] = defaultdict(set)
+        """tracks keys with gradient-like buffers, covariance-like buffers, etc for reprojecting"""
+
+        self._fake_params: dict[str, list[torch.Tensor]] = {}
+        """fake parameters for state keys and shape inference, key is name of child, value is list of fake parameters"""
+
 
     def set_param_groups(self, param_groups: Params):
         """Set custom parameter groups with per-parameter settings that this module will use."""
@@ -123,7 +130,9 @@ class Module(ABC):
             clone (bool):
                 If ``key`` exists, whether to clone ``tensors`` to avoid modifying buffers in-place.
                 If ``key`` doesn't exist, ``tensors`` are always returned without cloning
-            params (Iterable[torch.Tensor] | None, optional): pass None if ``tensors`` have different shape. Defaults to None.
+            params (Iterable[torch.Tensor] | None, optional):
+                pass None if ``tensors`` have different shape, it will create fake params from tensors
+                for state keys and shape inference. Defaults to None.
             grads (Sequence[torch.Tensor] | None, optional): grads. Defaults to None.
             loss (torch.Tensor | None, optional): loss. Defaults to None.
             closure (Callable | None, optional): closure. Defaults to None.
@@ -137,9 +146,26 @@ class Module(ABC):
             return tensors
 
         if clone: tensors = [t.clone() for t in tensors]
-        return step_tensors(modules=child, tensors=tensors, params=params, grads=grads,
+
+        # set fake params to same storage as tensors so as to not use any extra memory
+        # while they still refer to same python objects, so they can be used
+        # as state keys and for shape inference when params aren't given.
+        fake = params is None
+        if fake:
+            if key not in self._fake_params:
+                self._fake_params[key] = [torch.empty_like(t) for t in tensors]
+            params = self._fake_params[key]
+            _set_fake_params_(params, tensors)
+
+        update = step_tensors(modules=child, tensors=tensors, params=params, grads=grads,
                             loss=loss, closure=closure, objective=objective)
 
+        # set fake params storage to empty
+        if fake:
+            _empty_fake_param_storage_(params)
+
+        return update
+
 
     def __repr__(self):
         s = self.__class__.__name__
@@ -322,6 +348,48 @@ class Module(ABC):
         self.global_state[key] = value
         return value
 
+    def get_child_projected_buffers(self, key: str, buff: ProjectedBuffer | Sequence[ProjectedBuffer], params:Sequence[torch.Tensor] | None = None) -> list[list[torch.Tensor]]:
+        """if params is None, assumes fake parameters"""
+        if isinstance(buff, str): buff = (buff, )
+
+        child = self.children[key]
+        child.on_get_projected_buffers()
+        if params is None:
+            params = self._fake_params[key]
+
+        vals = []
+        for b in buff:
+            for buff_key in child._projected_keys[b]:
+                state = child.state[params[0]]
+                if buff_key in state:
+                    tensors = [child.state[p][buff_key] for p in params]
+                    if isinstance(tensors[0], torch.Tensor):
+                        vals.append(tensors)
+                    else: # its usually a deque
+                        assert isinstance(tensors[0], Sequence), type(tensors[0])
+                        vals.extend(zip(*tensors))
+
+                elif buff_key in child.global_state:
+                    val = child.global_state[buff_key]
+                    if len(val) == 0: continue
+                    if isinstance(val[0], torch.Tensor):
+                        vals.append(val)
+                    else:
+                        assert isinstance(val[0], Sequence)
+                        vals.extend(zip(*vals))
+
+        # recursively do this on children,
+        # note that if params are fake, children will have same fake params
+        # unless that child steps with something else. I don't think that is feasible to support it
+        for c in child.children:
+            vals.extend(child.get_child_projected_buffers(c, buff, params=params))
+
+        return vals
+
+    def add_projected_keys(self, buffer: ProjectedBuffer, *keys):
+        for k in keys: self._projected_keys[buffer].add(k)
+
+
     # ---------------------------- OVERRIDABLE METHODS --------------------------- #
     def update(self, objective:"Objective") -> None:
         """Updates internal state of this module. This should not modify ``objective.update``.
@@ -394,6 +462,9 @@ class Module(ABC):
         """
         for c in self.children.values(): c.reset_for_online()
 
+    def on_get_projected_buffers(self):
+        """runs before projected buffers are accessed"""
+
     def _extra_pack(self) -> dict:
         """extra information to store in ``state_dict`` of this optimizer.
         Will be passed to ``_extra_unpack`` when loading the ``state_dict``."""

torchzero/core/transform.py

@@ -1,12 +1,12 @@
 from abc import ABC, abstractmethod
 from collections.abc import Mapping, Sequence
 from operator import itemgetter
-from typing import Any, final, cast, TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, cast, final
 
 import torch
 
+from ..utils import safe_dict_update_, vec_to_tensors
 from .module import Module
-from ..utils import vec_to_tensors, safe_dict_update_
 
 if TYPE_CHECKING:
     from .chain import Chainable
@@ -31,7 +31,7 @@ class Transform(Module):
 
         self._objective = None
         if inner is not None:
-            self.set_child("inner", inner)
+            self.set_child("__inner", inner)
 
     # settings shouldn't mutate, so they are typed as Sequence[Mapping]
     def update_states(self, objective: "Objective", states: list[dict[str, Any]], settings: Sequence[Mapping[str, Any]]) -> None:
@@ -70,8 +70,8 @@ class Transform(Module):
     def apply(self, objective: "Objective"):
 
         # inner step
-        if "inner" in self.children:
-            inner = self.children["inner"]
+        if "__inner" in self.children:
+            inner = self.children["__inner"]
             objective = inner.step(objective)
 
         # apply and return
@@ -128,6 +128,7 @@ class TensorTransform(Transform):
         self._uses_grad = uses_grad
         self._uses_loss = uses_loss
 
+
     # ------------------------------- single tensor ------------------------------ #
     def single_tensor_initialize(
         self,

torchzero/linalg/eigh.py

@@ -10,30 +10,18 @@ from .svd import tall_reduced_svd_via_eigh
 
 # https://arxiv.org/pdf/2110.02820
 def nystrom_approximation(
-    A_mv: Callable[[torch.Tensor], torch.Tensor] | None,
-    A_mm: Callable[[torch.Tensor], torch.Tensor] | None,
-    ndim: int,
-    rank: int,
-    device,
-    orthogonalize_method: OrthogonalizeMethod = 'qr',
+    Omega: torch.Tensor,
+    AOmega: torch.Tensor,
     eigv_tol: float = 0,
-    dtype = torch.float32,
-    generator = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Computes Nyström approximation to positive-semidefinite A factored as Q L Q^T (truncatd eigenvalue decomp),
     returns ``(L, Q)``.
 
     A is ``(m,m)``, then Q is ``(m, rank)``; L is a ``(rank, )`` vector - diagonal of ``(rank, rank)``"""
-    # basis
-    O = torch.randn((ndim, rank), device=device, dtype=dtype, generator=generator) # Gaussian test matrix
-    O = orthogonalize(O, method=orthogonalize_method) # Thin QR decomposition # pylint:disable=not-callable
-
-    # Y = AΩ
-    AO = mm(A_mv=A_mv, A_mm=A_mm, X=O)
 
-    v = torch.finfo(dtype).eps * torch.linalg.matrix_norm(AO, ord='fro') # Compute shift # pylint:disable=not-callable
-    Yv = AO + v*O # Shift for stability
-    C = torch.linalg.cholesky_ex(O.mT @ Yv)[0] # pylint:disable=not-callable
+    v = torch.finfo(AOmega.dtype).eps * torch.linalg.matrix_norm(AOmega, ord='fro') # Compute shift # pylint:disable=not-callable
+    Yv = AOmega + v*Omega # Shift for stability
+    C = torch.linalg.cholesky_ex(Omega.mT @ Yv)[0] # pylint:disable=not-callable
     B = torch.linalg.solve_triangular(C, Yv.mT, upper=False, unitriangular=False).mT # pylint:disable=not-callable
 
     # Q, S, _ = torch_linalg.svd(B, full_matrices=False) # pylint:disable=not-callable
@@ -138,26 +126,35 @@ def eigh_plus_uuT(
 
     return L_prime, Q_prime
 
-def eigh_plus_UUT(
+def eigh_plus_UUt(
     L: torch.Tensor,
     Q: torch.Tensor,
     U: torch.Tensor,
-    alpha: float = 1,
+    alpha: float | torch.Tensor = 1,
     tol = None,
-    retry_float64: bool = False,
-):
+    ortho_method: OrthogonalizeMethod = 'qr',
+    retry_float64=True,
+) -> tuple[torch.Tensor, torch.Tensor] | tuple[None, None]:
     """
     compute eigendecomposition of Q L Q^T + alpha * (U U^T), where Q is ``(m, rank)`` and L is ``(rank, )``,
     U is ``(m, k)`` where k is rank of correction
+
+    returns ``(L, Q)``
     """
     if U.size(1) == 1:
-        return eigh_plus_uuT(L, Q, U[:,0], alpha=alpha, tol=tol, retry_float64=retry_float64)
+        return eigh_plus_uuT(L, Q, U[:,0], alpha=float(alpha), tol=tol)
+
+    # make alpha shape (k, )
+    k = U.size(1)
+    if isinstance(alpha, torch.Tensor):
+        alpha = torch.broadcast_to(alpha, (k, ))
+    else:
+        alpha = torch.full((k,), float(alpha), device=U.device, dtype=U.dtype)
 
     if tol is None: tol = torch.finfo(Q.dtype).eps
     m, r = Q.shape
-
-    Z = Q.T @ U  # (r, k)
-    U_res = U - Q @ Z  # (m, k)
+    QtU = Q.T @ U  # (r, k)
+    U_res = U - Q @ QtU  # (m, k)
 
     # find cols of U not in col space of Q
     res_norms = torch.linalg.vector_norm(U_res, dim=0) # pylint:disable=not-callable
@@ -167,23 +164,26 @@ def eigh_plus_UUT(
     if k_prime == 0:
         # all cols are in Q
         B = Q
-        C = Z # (r x k)
+        C = QtU # (r x k)
         r_new = r
     else:
         # orthonormalize directions that aren't in Q
         U_new = U_res[:, new_indices]
-        Q_u, _ = torch_linalg.qr(U_new, mode='reduced', retry_float64=retry_float64)
+        Q_u = orthogonalize(U_new, method=ortho_method)
         B = torch.hstack([Q, Q_u])
-        C = torch.vstack([Z, Q_u.T @ U])
+        C = torch.vstack([QtU, Q_u.T @ U_res])
         r_new = r + k_prime
 
-
     # project and compute new eigendecomposition
     A_proj = torch.zeros((r_new, r_new), device=Q.device, dtype=Q.dtype)
     A_proj[:r, :r] = L.diag_embed()
-    A_proj.addmm_(C, C.T, alpha=alpha)
+    # A_proj += (C @ C.T).mul_(alpha)
+    A_proj.addmm_(C * alpha, C.T)
 
-    L_prime, S = torch_linalg.eigh(A_proj, retry_float64=retry_float64)
+    try:
+        L_prime, S = torch_linalg.eigh(A_proj, retry_float64=retry_float64)
+    except torch.linalg.LinAlgError:
+        return None, None
 
     # unproject and sort
     Q_prime = B @ S
@@ -194,60 +194,108 @@ def eigh_plus_UUT(
     return L_prime, Q_prime
 
 
-def eigh_plus_UVT_symmetrize(
-    Q: torch.Tensor,
+def eigh_plus_UUt_mm(
+    # A1 = Q @ diag(L) @ Q.T
     L: torch.Tensor,
+    Q: torch.Tensor,
+
+    # A2 = U @ U.T
     U: torch.Tensor,
-    V: torch.Tensor,
-    alpha: float,
-    retry_float64: bool = False,
 
-):
+    # rhs
+    B: torch.Tensor,
+
+    # weights
+    w1: float,
+    w2: float | torch.Tensor,
+
+) -> torch.Tensor:
     """
-    Q is ``(m, rank)``; L is ``(rank, )``; U and V are the low rank correction such that U V^T is ``(m, m)``.
+    Computes ``(w1 * (Q L Q^T) + (U diag(w2) U^T) @ B``,
 
-    This computes eigendecomposition of A, where
+    Q is ``(m, rank)``, L is ``(rank, rank)``, U is ``(m, z)``, B is ``(m, k)``.
 
-    ``M = Q diag(L) Q^T + alpha * (U V^T)``;
+    Returns ``(m, k)``
+    """
+    # sketch Q L Q^T
+    QtB = Q.T @ B # (rank, k)
+    LQtB = L.unsqueeze(1) * QtB  # (rank, k)
+    sketch1 = Q @ LQtB  # (m, k)
+
+    # skecth U U^T
+    UtB = U.T @ B # (z, k)
+    if isinstance(w2, torch.Tensor) and w2.numel() > 1: w2UtB = w2.unsqueeze(-1) * UtB
+    else:  w2UtB = w2 * UtB
+    sketch2 = U @ w2UtB # (m, k)
+
+    return w1 * sketch1 + sketch2
 
-    ``A = (M + M^T) / 2``
+
+def randomized_eigh_plus_UUt(
+    L1: torch.Tensor,
+    Q1: torch.Tensor,
+    U: torch.Tensor,
+    w1: float,
+    w2: float | torch.Tensor,
+    oversampling_p: int,
+    rank: int,
+    eig_tol: float,
+    damping: float,
+    rdamping: float,
+    ortho_method: OrthogonalizeMethod = 'qr',
+) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+    """
+    compute randomized eigendecomposition of w1 * Q L Q^T + w2 * (U U^T),
+    where Q is ``(m, rank)`` and L is ``(rank, )``,
+    U is ``(m, k)`` where k is rank of correction, returns ``(L, Q)``
     """
-    m, rank = Q.shape
-    _, k = V.shape
+    n = Q1.shape[0]
+    device = Q1.device
+    dtype = Q1.dtype
+    l = rank + oversampling_p
 
-    # project U and V out of the Q subspace via Gram-schmidt
-    Q_T_U = Q.T @ U
-    U_perp = U - Q @ Q_T_U
+    # gaussian test matrix
+    Omega = torch.randn(n, l, device=device, dtype=dtype)
 
-    Q_T_V = Q.T @ V
-    V_perp = V - Q @ Q_T_V
+    # sketch
+    AOmega = eigh_plus_UUt_mm(L1, Q1, U, Omega, w1, w2)
+    Q = orthogonalize(AOmega, ortho_method)
 
-    R = torch.hstack([U_perp, V_perp])
-    Q_perp, _ = torch_linalg.qr(R, retry_float64=retry_float64)
+    AQ = eigh_plus_UUt_mm(L1, Q1, U, Q, w1, w2)
+    QtAQ = Q.T @ AQ
 
-    Q_B = torch.hstack([Q, Q_perp])
-    r_B = Q_B.shape[1]
+    W = (QtAQ + QtAQ.T) / 2.0
 
-    # project, symmetrize and compute new eigendecomposition
-    A_proj = torch.zeros((r_B, r_B), device=Q.device, dtype=Q.dtype)
-    A_proj[:rank, :rank] = L.diag_embed()
+    # compute new L and Q
+    try:
+        L_prime, S = torch.linalg.eigh(W) # pylint:disable=not-callable
+    except torch.linalg.LinAlgError:
+        return L1, Q1
 
-    Q_perp_T_U = Q_perp.T @ U
-    Q_B_T_U = torch.vstack([Q_T_U, Q_perp_T_U])
+    L_prime, S = regularize_eigh(L=L_prime, Q=S, truncate=rank, tol=eig_tol, damping=damping, rdamping=rdamping)
 
-    Q_perp_T_V = Q_perp.T @ V
-    Q_B_T_V = torch.vstack([Q_T_V, Q_perp_T_V])
+    if L_prime is None or S is None:
+        return L1, Q1
 
-    update_proj = Q_B_T_U @ Q_B_T_V.T + Q_B_T_V @ Q_B_T_U.T
-    A_proj.add_(update_proj, alpha=alpha/2)
+    return L_prime, Q @ S
 
-    L_prime, S = torch_linalg.eigh(A_proj, retry_float64=retry_float64)
 
-    # unproject and sort
-    Q_prime = Q_B @ S
+def rank1_eigh(v: torch.Tensor):
+    """returns ``(L, Q)`` of ``(v v^T)``"""
+    vv = v.dot(v)
+    norm = vv.sqrt().clip(min=torch.finfo(vv.dtype).tiny * 2)
 
-    idx = torch.argsort(L_prime)
-    L_prime = L_prime[idx]
-    Q_prime = Q_prime[:, idx]
+    L = vv.unsqueeze(0) # (rank, )
+    Q = v.unsqueeze(-1) / norm # (m, rank)
 
-    return L_prime, Q_prime
+    return L, Q
+
+def low_rank_eigh(U: torch.Tensor):
+    """returns ``(L, Q)`` of ``alpha * (U U^T)`` (from GGT)"""
+    M = U.T @ U
+    L, S = torch.linalg.eigh(M) # pylint:disable=not-callable
+
+    Q = U @ S
+    Q /= torch.sqrt(L).clip(min=torch.finfo(L.dtype).tiny * 2)
+
+    return L, Q

torchzero/linalg/linear_operator.py

@@ -425,3 +425,4 @@ class Eigendecomposition(LinearOperator):
     def size(self):
         n = self.Q.size(0)
         return (n,n)
+

torchzero/linalg/orthogonalize.py

@@ -51,9 +51,6 @@ def zeropower_via_newtonschulz5(G: torch.Tensor, coeffs=_NS_COEFFS) -> torch.Ten
     return X.to(G.dtype)
 
 def zeropower_via_svd(A: torch.Tensor) -> torch.Tensor:
-    """
-    Applies to first 2 dims and isn't batched - rest of dimensions are flattened.
-    """
     try:
         U, S, Vt = torch_linalg.svd(A, full_matrices=False, retry_float64=True) # pylint:disable=not-callable
     except torch.linalg.LinAlgError:
@@ -84,9 +81,67 @@ def orthogonalize_via_qr(A: torch.Tensor):
 
     return Q
 
-OrthogonalizeMethod = Literal["newtonschulz", "svd", "qr"]
+# CODE FROM https://github.com/HomebrewML/HeavyBall/blob/main/heavyball/utils.py:
+
+## Based on https://arxiv.org/pdf/2505.16932v3
+# and https://github.com/NoahAmsel/PolarExpress/blob/5454910920ca8c65afda28820cdf9e49b9436ed0/polar_express.py#L69-L82
+# and https://github.com/thinking-machines-lab/manifolds/blob/89dcae50f01af59f1e0570289474da3a2ecaa60b/src/msign.py#L47
+#
+# under the MIT License
+# Coefficients are from https://arxiv.org/pdf/2505.16932v3
+ABC_LIST: list[tuple[float, float, float]] = [
+    (8.28721201814563, -23.595886519098837, 17.300387312530933),
+    (4.107059111542203, -2.9478499167379106, 0.5448431082926601),
+    (3.9486908534822946, -2.908902115962949, 0.5518191394370137),
+    (3.3184196573706015, -2.488488024314874, 0.51004894012372),
+    (2.300652019954817, -1.6689039845747493, 0.4188073119525673),
+    (1.891301407787398, -1.2679958271945868, 0.37680408948524835),
+    (1.8750014808534479, -1.2500016453999487, 0.3750001645474248),
+    (1.875, -1.25, 0.375),
+]
+
+# safety factor for numerical stability (but exclude last polynomial)
+ABC_LIST_STABLE: list[tuple[float, float, float]] = [
+    (a / 1.01, b / 1.01**3, c / 1.01**5) for (a, b, c) in ABC_LIST[:-1]
+] + [ABC_LIST[-1]]
+
+
+def msign(G: torch.Tensor, steps: int = 10, eps: float = 1e-7) -> torch.Tensor:
+    """
+    Polar Express algorithm for the matrix sign function:
+    https://arxiv.org/abs/2505.16932
+    """
+    assert G.ndim >= 2
+    should_transpose: bool = G.size(-2) > G.size(-1)
+
+    x = G
+    if should_transpose:
+        x = x.mT
+
+    x = x / (x.norm(dim=(-2, -1), keepdim=True) * 1.01 + eps)
+
+    for step in range(steps):
+        a, b, c = ABC_LIST_STABLE[step] if step < len(ABC_LIST_STABLE) else ABC_LIST_STABLE[-1]
+        s = x @ x.mT
+        # goal is to compute x = a x + b S x + c S^2 x
+        # we can break this up into: x = (a I + (b I + c S) S) x
+        y = c * s
+        y.diagonal(dim1=-2, dim2=-1).add_(b)
+        y = y @ s
+        y.diagonal(dim1=-2, dim2=-1).add_(a)
+        x = y @ x
+
+    if should_transpose:
+        x = x.mT
+    return x.float()
+
+
+###### END
+
+OrthogonalizeMethod = Literal["newtonschulz", "ns5", "polar_express", "svd", "qr", "eigh"]
 def orthogonalize(A: torch.Tensor, method: OrthogonalizeMethod) -> torch.Tensor:
-    if method == "newtonschulz": return zeropower_via_newtonschulz5(A)
+    if method in ("newtonschulz", "ns5"): return zeropower_via_newtonschulz5(A)
+    if method == "polar_express": return msign(A)
     if method == "svd": return zeropower_via_svd(A)
     if method == "qr": return orthogonalize_via_qr(A)
     if method == "eigh": return zeropower_via_eigh(A)

torchzero/linalg/sketch.py

@@ -0,0 +1,39 @@
+import math
+
+import torch
+
+from .orthogonalize import orthogonalize_via_qr
+from .linear_operator import LinearOperator, Dense
+
+
+class Permutation(LinearOperator):
+    def __init__(self, indices:torch.Tensor):
+        self.indices = indices
+        self.device = indices.device
+
+    def matvec(self, x):
+        return x[self.indices]
+
+    def matmat(self, X):
+        return Dense(X[:, self.indices])
+
+def orthonormal_sketch(m, k, dtype, device, generator):
+    return orthogonalize_via_qr(torch.randn(m, k, dtype=dtype, device=device, generator=generator))
+
+def rademacher_sketch(m, k, dtype, device, generator):
+    rademacher = torch.bernoulli(torch.full((m, k), 0.5, device=device, dtype=dtype), generator = generator).mul_(2).sub_(1)
+    return rademacher.mul_(1 / math.sqrt(m))
+
+def row_sketch(m, k, dtype, device, generator):
+    weights = torch.ones(m, dtype=dtype, device=device)
+    indices = torch.multinomial(weights, k, replacement=False, generator=generator)
+
+    P = torch.zeros(m, k, dtype=dtype, device=device)
+    P[indices, range(k)] = 1
+    return P
+
+def topk_rows_sketch(v: torch.Tensor, m, k, dtype, device):
+    _, indices = torch.topk(v, k)
+    P = torch.zeros(m, k, dtype=dtype, device=device)
+    P[indices, range(k)] = 1
+    return P

torchzero/modules/__init__.py

@@ -21,3 +21,4 @@ from .variance_reduction import *
 from .weight_decay import *
 from .wrappers import *
 from .zeroth_order import *
+from .basis import *

torchzero/modules/adaptive/adagrad.py

@@ -40,6 +40,7 @@ class Adagrad(TensorTransform):
         super().__init__(defaults=defaults, inner=inner)
 
         self.set_child('accumulator', accumulator_tfm)
+        self.add_projected_keys("grad", "accumulator")
 
     @torch.no_grad
     def single_tensor_initialize(self, tensor, param, grad, loss, state, setting):
@@ -235,6 +236,7 @@ class FullMatrixAdagrad(TensorTransform):
         super().__init__(defaults=defaults, inner=inner, concat_params=concat_params)
 
         self.set_child("accumulator", accumulator_tfm)
+        self.add_projected_keys("covariance", "accumulator")
 
     @torch.no_grad
     def single_tensor_update(self, tensor, param, grad, loss, state, setting):

torchzero/modules/adaptive/adam.py

@@ -38,6 +38,9 @@ class Adam(TensorTransform):
         self.set_child('exp_avg', exp_avg_tfm)
         self.set_child('exp_avg_sq', exp_avg_sq_tfm)
 
+        self.add_projected_keys("grad", "exp_avg")
+        self.add_projected_keys("grad_sq", "exp_avg_sq", "max_exp_avg_sq")
+
     @torch.no_grad
     def multi_tensor_update(self, tensors, params, grads, loss, states, settings):
         self.increment_counter("step", start=0)
@@ -81,4 +84,5 @@ class Adam(TensorTransform):
             exp_avg = exp_avg * alpha
 
         # ---------------------------------- update ---------------------------------- #
-        return exp_avg / exp_avg_sq.sqrt().add_(eps)
+        return exp_avg / exp_avg_sq.sqrt().add_(eps)
+

torchzero/modules/adaptive/adan.py

@@ -87,6 +87,9 @@ class Adan(TensorTransform):
         self.set_child("v", v_tfm)
         self.set_child("n", n_tfm)
 
+        self.add_projected_keys("grad_sq", "m", "v", "g_prev")
+        self.add_projected_keys("grad", "n")
+
     @torch.no_grad
     def multi_tensor_update(self, tensors, params, grads, loss, states, settings):
         tensors = TensorList(tensors)