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

torchzero/modules/experimental/adasoap.py

@@ -2,11 +2,18 @@ from operator import itemgetter
 
 import torch
 
-from ...core import Chainable, Transform, apply
+from ...core import Chainable, Transform
 from ...modules.optimizers.shampoo import _merge_small_dims, _unmerge_small_dims
+from ..optimizers.soap import (
+    get_orthogonal_matrix,
+    get_orthogonal_matrix_QR,
+    project,
+    project_back,
+)
+
 
 @torch.no_grad
-def update_soap_covariances_(
+def update_adasoap_covariances_(
     grad: torch.Tensor,
     GGs_: list[torch.Tensor | None],
     GG_sqs: list[torch.Tensor | None],
@@ -24,127 +31,16 @@ def update_soap_covariances_(
         if beta is None: GG.add_(torch.tensordot(grad, grad, (axes, axes))) # pyright:ignore[reportArgumentType]
         else: GG.lerp_(torch.tensordot(grad, grad, (axes, axes)), 1-beta) # pyright:ignore[reportArgumentType]
 
-@torch.no_grad
-def project(tensors: torch.Tensor, Q: list[torch.Tensor | None]):
-    """
-    Projects the gradient to the eigenbases of the preconditioner.
-    """
-    for mat in Q:
-        if mat is None: continue
-        if len(mat) > 0:
-            tensors = torch.tensordot(tensors, mat, dims=[[0], [0]]) # pyright:ignore[reportArgumentType]
-        else:
-            # I don't understand this part but it is in https://github.com/nikhilvyas/SOAP/blob/main/soap.py
-            permute_order = list(range(1, len(tensors.shape))) + [0]
-            tensors = tensors.permute(permute_order)
-
-    return tensors
-
-@torch.no_grad
-def project_back(tensors: torch.Tensor, Q: list[torch.Tensor| None]):
-    """
-    Projects the gradient back to the original space.
-    """
-    for mat in Q:
-        if mat is None: continue
-        if len(mat) > 0:
-            tensors = torch.tensordot(tensors, mat,dims=[[0], [1]]) # pyright:ignore[reportArgumentType]
-        else:
-            permute_order = list(range(1, len(tensors.shape))) + [0]
-            tensors = tensors.permute(permute_order)
-
-    return tensors
-
-# function from https://github.com/nikhilvyas/SOAP/blob/main/soap.py
-@torch.no_grad
-def get_orthogonal_matrix(mat: list[torch.Tensor | None]):
-    """
-    Computes the eigenbases of the preconditioner using torch.linalg.eigh decomposition.
-    """
-    matrix = []
-    float_data = False
-    original_type = original_device = None
-    for m in mat:
-        if m is None: continue
-        if len(m) == 0:
-            matrix.append([])
-            continue
-        if m.dtype != torch.float:
-            original_type = m.dtype
-            original_device = m.device
-            matrix.append(m.float())
-        else:
-            float_data = True
-            matrix.append(m)
-
-    final = []
-    for m in matrix:
-        if len(m) == 0:
-            final.append([])
-            continue
-        try:
-            _, Q = torch.linalg.eigh(m+1e-30*torch.eye(m.shape[0], device=m.device)) # pylint:disable=not-callable
-        except Exception:
-            _, Q = torch.linalg.eigh(m.to(torch.float64)+1e-30*torch.eye(m.shape[0], device=m.device)) # pylint:disable=not-callable
-            Q = Q.to(m.dtype)
-        Q = torch.flip(Q, [1])
-
-        if not float_data:
-            Q = Q.to(original_device).type(original_type)
-        final.append(Q)
-    return final
-
-# function from https://github.com/nikhilvyas/SOAP/blob/main/soap.py#L240
-@torch.no_grad
-def get_orthogonal_matrix_QR(exp_avg_sq: torch.Tensor, GG: list[torch.Tensor | None], Q_list: list[torch.Tensor | None]):
-    """
-    Computes the eigenbases of the preconditioner using one round of power iteration
-    followed by torch.linalg.qr decomposition.
-    """
-    matrix = []
-    orth_matrix = []
-    float_data = False
-    original_type = original_device = None
-    for m,o in zip(GG, Q_list):
-        if m is None: continue
-        assert o is not None
-
-        if len(m) == 0:
-            matrix.append([])
-            orth_matrix.append([])
-            continue
-        if m.data.dtype != torch.float:
-            original_type = m.data.dtype
-            original_device = m.data.device
-            matrix.append(m.data.float())
-            orth_matrix.append(o.data.float())
-        else:
-            float_data = True
-            matrix.append(m.data.float())
-            orth_matrix.append(o.data.float())
-
-    final = []
-    for ind, (m,o) in enumerate(zip(matrix, orth_matrix)):
-        if len(m)==0:
-            final.append([])
-            continue
-        est_eig = torch.diag(o.T @ m @ o)
-        sort_idx = torch.argsort(est_eig, descending=True)
-        exp_avg_sq = exp_avg_sq.index_select(ind, sort_idx)
-        o = o[:,sort_idx]
-        power_iter = m @ o
-        Q, _ = torch.linalg.qr(power_iter) # pylint:disable=not-callable
-
-        if not float_data:
-            Q = Q.to(original_device).type(original_type)
-        final.append(Q)
-
-    return final, exp_avg_sq
 
 class AdaSOAP(Transform):
-    """SOAP with diagonally preconditioned GG^Ts
+    """SOAP with diagonally preconditioned GG^Ts.
+
+    .. warning::
+        Experimental.
 
     precond_beta - beta for GG^T squares
+
+    Verdict: It works, but it is about the same performance as Adam, but maybe more tuning potential?
     """
     def __init__(
         self,
@@ -180,15 +76,14 @@ class AdaSOAP(Transform):
         super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         updates = []
         # update preconditioners
-        for i,(p,t) in enumerate(zip(params, tensors)):
-            state = self.state[p]
-            settings = self.settings[p]
+        for i,(p,t, state, setting) in enumerate(zip(params, tensors, states, settings)):
+
             beta1, beta2, shampoo_beta, merge_small, max_dim, precondition_1d, eps, unprojected_exp_avg,alpha = itemgetter(
-                'beta1', 'beta2', 'shampoo_beta', 'merge_small', 'max_dim', 'precondition_1d', 'eps', 'unprojected_exp_avg','alpha')(settings)
-            precond_beta = settings['precond_beta']
+                'beta1', 'beta2', 'shampoo_beta', 'merge_small', 'max_dim', 'precondition_1d', 'eps', 'unprojected_exp_avg','alpha')(setting)
+            precond_beta = setting['precond_beta']
 
             if merge_small:
                 t, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(t, max_dim)
@@ -213,7 +108,7 @@ class AdaSOAP(Transform):
 
                 if state['GG'] is not None:
                     assert state['GG_sq'] is not None
-                    update_soap_covariances_(t, GGs_=state['GG'], GG_sqs=state['GG_sq'], beta=shampoo_beta, precond_beta=precond_beta)
+                    update_adasoap_covariances_(t, GGs_=state['GG'], GG_sqs=state['GG_sq'], beta=shampoo_beta, precond_beta=precond_beta)
                     GG_precond = [GG / (GG_sq+1e-8) if GG is not None and GG_sq is not None else None for GG, GG_sq in zip(state['GG'], state['GG_sq'])]
                     state['Q'] = get_orthogonal_matrix(GG_precond)
 
@@ -259,7 +154,7 @@ class AdaSOAP(Transform):
             if t_projected is not None:
                 update = project_back(update, state["Q"])
 
-            if settings['bias_correction']:
+            if setting['bias_correction']:
                 bias_correction1 = 1.0 - beta1 ** (state["step"]+1)
                 bias_correction2 = 1.0 - beta2 ** (state["step"]+1)
                 update *= ((bias_correction2 ** .5) / bias_correction1) * alpha
@@ -274,9 +169,9 @@ class AdaSOAP(Transform):
 
             # Update is done after the gradient step to avoid using current gradients in the projection.
             if state['GG'] is not None:
-                update_soap_covariances_(t, GGs_=state['GG'], GG_sqs=state['GG_sq'], beta=shampoo_beta, precond_beta=precond_beta)
+                update_adasoap_covariances_(t, GGs_=state['GG'], GG_sqs=state['GG_sq'], beta=shampoo_beta, precond_beta=precond_beta)
                 GG_precond = [GG / (GG_sq+1e-8) if GG is not None and GG_sq is not None else None for GG, GG_sq in zip(state['GG'], state['GG_sq'])]
-                if state['step'] % settings['precond_freq'] == 0:
+                if state['step'] % setting['precond_freq'] == 0:
                     state['Q'], state['exp_avg_sq'] = get_orthogonal_matrix_QR(exp_avg_sq, GG_precond, state['Q'])
 
         return updates

torchzero/modules/experimental/cosine.py

@@ -0,0 +1,214 @@
+"""A bunch of useless modules that I hate and that didn't work"""
+import torch
+
+from ...core import Chainable, Transform, apply_transform
+from ...utils import NumberList, TensorList, as_tensorlist, unpack_dicts, unpack_states
+
+
+class CosineStepSize(Transform):
+    """Adaptive step size based on cosine similarity
+
+    VERDICT: Useless. This is too unstable.
+
+    Args:
+        scale (float, optional): cosine similarity multiplier. Defaults to 0.95.
+        init (float, optional): initial step size. Defaults to 1.
+        eps (float, optional): epsilon for division stability. Defaults to 1e-12.
+        target_cossim (float, optional): cosine similarity needs to be above this to increase step size. Defaults to 1e-8.
+        inner (Chainable | None, optional):
+            inner modules applied after calculating cosine similarity and before step size correction. Defaults to None.
+    """
+    def __init__(self, scale:float = 0.95, init:float=1, eps:float=1e-12, inner:Chainable | None = None):
+        defaults = dict(scale=scale, init=init, eps=eps)
+        super().__init__(defaults, uses_grad=False)
+        if inner is not None: self.set_child('inner', inner)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        scale, init = unpack_dicts(settings, 'scale', 'init', cls=NumberList)
+        unpack_states(states, tensors, 'alpha', init=init, cls=NumberList) # initializes alpha to init
+        eps = settings[0]['eps']
+
+        tensors = as_tensorlist(tensors)
+        prev = unpack_states(states, tensors, 'prev', init=tensors, cls=TensorList)
+
+        tensors_norm = tensors.global_vector_norm()
+        cos_sim = (tensors.dot(prev) / (tensors_norm * prev.global_vector_norm()).clip(min=eps)).item()
+
+        if 'inner' in self.children:
+            tensors = as_tensorlist(apply_transform(self.children['inner'], tensors, params, grads, loss))
+
+        new_alpha = []
+        for s, sc in zip(states, scale):
+            s['alpha'] *= 1 + cos_sim * sc
+            new_alpha.append(s['alpha'])
+
+        tensors.mul_(new_alpha)
+        prev.copy_(tensors)
+
+        return tensors
+
+
+
+class CosineDebounce(Transform):
+    """Debouncing when cosine similarity is less than 0.
+
+    VERDICT: Useless. This doesn't help at all.
+
+    Args:
+        scale (float, optional): cosine similarity multiplier. Defaults to 0.95.
+        eps (float, optional): epsilon for division stability. Defaults to 1e-12.
+        inner (Chainable | None, optional):
+            inner modules applied after calculating cosine similarity and before debouncing correction. Defaults to None.
+    """
+    def __init__(self, scale:float = 0.95, eps:float=1e-12, damping:float=0.95, inner:Chainable | None = None):
+        defaults = dict(scale=scale, eps=eps, damping=damping)
+        super().__init__(defaults, uses_grad=False)
+        if inner is not None: self.set_child('inner', inner)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        scale, damping = unpack_dicts(settings, 'scale', 'damping', cls=NumberList)
+        eps = settings[0]['eps']
+
+        tensors = as_tensorlist(tensors)
+        prev = unpack_states(states, tensors, 'prev', init=tensors, cls=TensorList).mul_(damping)
+
+        tensors_norm = tensors.global_vector_norm()
+        cos_sim = (tensors.dot(prev) / (tensors_norm * prev.global_vector_norm()).clip(min=eps)).item()
+
+        if 'inner' in self.children:
+            tensors = as_tensorlist(apply_transform(self.children['inner'], tensors, params, grads, loss))
+
+        if cos_sim < -eps:
+            undo = prev.neg().mul_(-cos_sim * scale)
+            comb = prev.graft(tensors).add_(tensors).graft_(prev).mul_(-cos_sim*scale)
+            tensors = undo.add_(comb)
+
+        prev.copy_(tensors)
+        return tensors
+
+
+
+class CosineMomentum(Transform):
+    """Beta depends on cosine similarity. At cossim=1, beta is 0. At cossim=-1, beta is 2^power. This basically removes oscillations.
+
+    VERDICT: Useless. Worse than all other momentums.
+
+    Args:
+        scale (float, optional): cosine similarity multiplier. Defaults to 1.
+        nesterov (float, optional): whether to use nesterov momentum. Defaults to False.
+        power (float, optional): power for beta. Defaults to 1.
+        eps (float, optional): epsilon for division stability. Defaults to 1e-12.
+        inner (Chainable | None, optional):
+            inner modules applied after calculating cosine similarity and before updating exponential moving average. Defaults to None.
+    """
+    def __init__(self, scale:float = 1, nesterov: bool = False, power: float = 1, eps:float=1e-12, inner:Chainable | None = None):
+        defaults = dict(scale=scale, eps=eps, nesterov=nesterov, power=power)
+        super().__init__(defaults, uses_grad=False)
+        if inner is not None: self.set_child('inner', inner)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        scale, power = unpack_dicts(settings, 'scale', 'power', cls=NumberList)
+        eps = settings[0]['eps']
+        nesterov = settings[0]['nesterov']
+        exp_avg = unpack_states(states, tensors, 'exp_avg', cls=TensorList)
+
+        tensors = as_tensorlist(tensors)
+
+        tensors_norm = tensors.global_vector_norm()
+        cos_sim = (tensors.dot(exp_avg) / (tensors_norm * exp_avg.global_vector_norm()).clip(min=eps)).item()
+
+        if 'inner' in self.children:
+            tensors = as_tensorlist(apply_transform(self.children['inner'], tensors, params, grads, loss))
+
+        beta = (1 - (cos_sim*scale)) ** power
+        if nesterov:
+            exp_avg.add_(tensors.mul(beta))
+            return tensors.add_(exp_avg)
+        else:
+            exp_avg.add_(tensors.mul_(beta))
+            return exp_avg.clone()
+
+
+class AdaptiveDifference(Transform):
+    """VERDICT: Useless. Doesn't help (sort of to be expected)."""
+    def __init__(self, inner:Chainable | None = None):
+        defaults = dict()
+        super().__init__(defaults, uses_grad=False)
+        if inner is not None: self.set_child('inner', inner)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        tensors = as_tensorlist(tensors)
+        prev = unpack_states(states, tensors, 'prev', init=tensors, cls=TensorList)
+
+        diff = tensors - prev.graft_(tensors)
+        prev.copy_(tensors)
+
+        if 'inner' in self.children:
+            tensors = as_tensorlist(apply_transform(self.children['inner'], tensors, params, grads, loss))
+
+        tensors.add_(diff.graft_(tensors))
+
+        return tensors
+
+class AdaptiveDifferenceEMA(Transform):
+    """VERDICT: better than non-EMA but still useless."""
+    def __init__(self, beta=0.99, inner:Chainable | None = None):
+        defaults = dict(beta=beta)
+        super().__init__(defaults, uses_grad=False)
+        if inner is not None: self.set_child('inner', inner)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        tensors = as_tensorlist(tensors)
+        beta = unpack_dicts(settings, 'beta', cls=NumberList)
+        prev, diff_exp_avg = unpack_states(states, tensors, 'prev', 'diff_exp_avg', init=[tensors,torch.zeros_like], cls=TensorList)
+
+        diff = (tensors - prev.graft_(tensors)).graft_(tensors)
+        diff_exp_avg.lerp_(diff, 1-beta)
+        prev.copy_(tensors)
+
+        if 'inner' in self.children:
+            tensors = as_tensorlist(apply_transform(self.children['inner'], tensors, params, grads, loss))
+
+        tensors.add_(diff_exp_avg.graft(tensors))
+
+        return tensors
+
+
+class ScaledAdaptiveDifference(Transform):
+    """VERDICT: Useless and doesn't help."""
+    def __init__(self, scale=0.95, damping:float=0.99, inner:Chainable | None = None):
+        defaults = dict(scale=scale, damping=damping)
+        super().__init__(defaults, uses_grad=False)
+        if inner is not None: self.set_child('inner', inner)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        tensors = as_tensorlist(tensors)
+        scale, damping = unpack_dicts(settings, 'scale', 'damping', cls=NumberList)
+        prev_tensors, prev_update = unpack_states(states, tensors, 'prev', 'prev_update', init=[tensors,tensors], cls=TensorList)
+
+        cos_sim = (tensors.dot(prev_update) / (tensors.global_vector_norm() * prev_update.global_vector_norm()).clip(min=1e-10)).item()
+
+        if 'inner' in self.children:
+            tensors = as_tensorlist(apply_transform(self.children['inner'], tensors, params, grads, loss))
+
+        if cos_sim > 0:
+            tensors.add_(prev_tensors*(cos_sim*scale))
+
+        else:
+            undo = prev_tensors.neg().mul_(-cos_sim*scale)
+            comb = prev_tensors.graft(tensors).add_(tensors).graft_(prev_tensors).mul_(-cos_sim*scale)
+            tensors = undo.add_(comb).graft_((tensors-prev_tensors).mul_(damping))
+
+        diff = tensors - prev_tensors.graft_(tensors)
+        prev_tensors.copy_(tensors)
+        diff.graft_(tensors)
+        tensors.add_(diff)
+        prev_update.copy_(tensors)
+
+        return tensors

torchzero/modules/experimental/cubic_adam.py

@@ -0,0 +1,97 @@
+import torch
+
+from ...core import Transform
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
+
+
+def signed_cbrt(x: TensorList) -> TensorList:
+    return x.sign() * x.abs().pow(1/3)
+
+def cubic_adam_(
+    tensors: TensorList,
+    exp_avg_: TensorList,
+    exp_avg_sq_: TensorList,
+    exp_avg_cu_: TensorList,
+    alpha: float | NumberList,
+    beta1: float | NumberList,
+    beta2: float | NumberList,
+    beta3: float | NumberList,
+    eps: float | NumberList,
+    debiased: bool,
+    step: int,
+):
+    exp_avg_.lerp_(tensors, 1-beta1)
+    exp_avg_sq_.lerp_(tensors**2, 1-beta2)
+    exp_avg_cu_.lerp_(tensors**3, 1-beta3)
+
+    if debiased:
+        m1 = exp_avg_ / (1 - beta1 ** step)
+        m2 = exp_avg_sq_ / (1 - beta2 ** step)
+        m3 = exp_avg_cu_ / (1 - beta3 ** step)
+    else:
+        m1, m2, m3 = exp_avg_, exp_avg_sq_, exp_avg_cu_
+
+    # adam minimizes ax^2 + bx
+    # we are going to minimize ax^3 + bx^2 + cx
+    A = signed_cbrt(m3)
+    B = m2.sqrt()
+    C = m1
+    discriminant = B.pow(2) - 4 * A * C
+
+    denom = 2 * A
+    root = discriminant.clamp(min=0).sqrt_()
+
+    x0 = (-B + root) / (denom + eps)
+    x1 = (-B - root) / (denom + eps)
+
+    f0 = (A/3)*x0**3 + (B/2)*x0**2 + C*x0
+    f1 = (A/3)*x1**3 + (B/2)*x1**2 + C*x1
+
+    x_star = x0.where(f0 < f1, x1)
+
+    adam = -C / (B + eps)
+    x_star = adam.where(discriminant < 0, x_star)
+
+    return x_star.mul_(-alpha)
+
+class CubicAdam(Transform):
+    """Adam which has 3rd momentum and minimizes a cubic polynomial.
+
+    VERDICT: can outperform Adam very slightly. Usually very similar performance.
+
+    .. warning::
+        Experimental.
+
+    """
+    def __init__(
+        self,
+        beta1: float = 0.9,
+        beta2: float = 0.99,
+        beta3: float = 0.99,
+        eps: float = 1e-8,
+        debiased:bool=True,
+        alpha: float = 1.,
+    ):
+        defaults=dict(beta1=beta1,beta2=beta2,beta3=beta3,eps=eps,debiased=debiased,alpha=alpha)
+        super().__init__(defaults, uses_grad=False)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
+
+        beta1,beta2,beta3,eps,alpha=unpack_dicts(settings, 'beta1','beta2','beta3','eps','alpha', cls=NumberList)
+        exp_avg, exp_avg_sq, exp_avg_cu = unpack_states(states, tensors, 'exp_avg', 'exp_avg_sq', 'exp_avg_cu', cls=TensorList)
+
+        return cubic_adam_(
+            tensors=TensorList(tensors),
+            exp_avg_=exp_avg,
+            exp_avg_sq_=exp_avg_sq,
+            exp_avg_cu_=exp_avg_cu,
+            alpha=alpha,
+            beta1=beta1,
+            beta2=beta2,
+            beta3=beta3,
+            eps=eps,
+            debiased=settings[0]['debiased'],
+            step=step,
+        )

torchzero/modules/experimental/curveball.py

@@ -2,7 +2,7 @@ from typing import Literal
 from collections.abc import Callable
 import torch
 
-from ...core import Module, Target, Transform, Chainable, apply
+from ...core import Module, Target, Transform, Chainable, apply_transform
 from ...utils import NumberList, TensorList, as_tensorlist
 from ...utils.derivatives import hvp, hvp_fd_forward, hvp_fd_central
 
@@ -47,27 +47,27 @@ class CurveBall(Module):
         if inner is not None: self.set_child('inner', inner)
 
     @torch.no_grad
-    def step(self, vars):
+    def step(self, var):
 
-        params = vars.params
+        params = var.params
         settings = self.settings[params[0]]
         hvp_method = settings['hvp_method']
         h = settings['h']
 
-        precond_lr, momentum, reg = self.get_settings('momentum', 'decay_rate', 'reg', params=params, cls=NumberList)
+        precond_lr, momentum, reg = self.get_settings(params, 'precond_lr', 'momentum', 'reg', cls=NumberList)
 
 
-        closure = vars.closure
+        closure = var.closure
         assert closure is not None
 
-        z, Hz = self.get_state('z', 'Hz', params=params, cls=TensorList)
+        z, Hz = self.get_state(params, 'z', 'Hz', cls=TensorList)
 
         if hvp_method == 'autograd':
-            grad = vars.get_grad(create_graph=True)
+            grad = var.get_grad(create_graph=True)
             Hvp = hvp(params, grad, z)
 
         elif hvp_method == 'forward':
-            loss, Hvp = hvp_fd_forward(closure, params, z, h=h, g_0=vars.get_grad(), normalize=True)
+            loss, Hvp = hvp_fd_forward(closure, params, z, h=h, g_0=var.get_grad(), normalize=True)
 
         elif hvp_method == 'central':
             loss, Hvp = hvp_fd_central(closure, params, z, h=h, normalize=True)
@@ -79,11 +79,11 @@ class CurveBall(Module):
         Hz.set_(Hvp + z*reg)
 
 
-        update = vars.get_update()
+        update = var.get_update()
         if 'inner' in self.children:
-            update = apply(self.children['inner'], update, params, grads=vars.grad, vars=vars)
+            update = apply_transform(self.children['inner'], update, params, grads=var.grad, var=var)
 
         z = curveball(TensorList(update), z, Hz, momentum=momentum, precond_lr=precond_lr)
-        vars.update = z.neg()
+        var.update = z.neg()
 
-        return vars
+        return var

torchzero/modules/{projections → experimental}/dct.py

@@ -1,13 +1,13 @@
 from typing import Literal
 import torch
 import torch_dct
-from .projection import Projection
+from ..projections import ProjectionBase
 from ...core import Chainable
 
 def reverse_dims(t:torch.Tensor):
     return t.permute(*reversed(range(t.ndim)))
 
-class DCTProjection(Projection):
+class DCTProjection(ProjectionBase):
     # norm description copied from pytorch docstring
     """Project update into Discrete Cosine Transform space, requires `torch_dct` library.
 
@@ -34,8 +34,8 @@ class DCTProjection(Projection):
         super().__init__(modules, project_update=project_update, project_params=project_params, project_grad=project_grad, defaults=defaults)
 
     @torch.no_grad
-    def project(self, tensors, vars, current):
-        settings = self.settings[vars.params[0]]
+    def project(self, tensors, params, grads, loss, states, settings, current):
+        settings = settings[0]
         dims = settings['dims']
         norm = settings['norm']
 
@@ -54,18 +54,18 @@ class DCTProjection(Projection):
         return projected
 
     @torch.no_grad
-    def unproject(self, tensors, vars, current):
-        settings = self.settings[vars.params[0]]
+    def unproject(self, projected_tensors, params, grads, loss, projected_states, projected_settings, current):
+        settings = projected_settings[0]
         dims = settings['dims']
         norm = settings['norm']
 
         unprojected = []
-        for u in tensors:
-            dim = min(u.ndim, dims)
+        for t in projected_tensors:
+            dim = min(t.ndim, dims)
 
-            if dim == 1: idct = torch_dct.idct(u, norm = norm)
-            elif dim == 2: idct = torch_dct.idct_2d(u, norm=norm)
-            elif dim == 3: idct = torch_dct.idct_3d(u, norm=norm)
+            if dim == 1: idct = torch_dct.idct(t, norm = norm)
+            elif dim == 2: idct = torch_dct.idct_2d(t, norm=norm)
+            elif dim == 3: idct = torch_dct.idct_3d(t, norm=norm)
             else: raise ValueError(f"Unsupported number of dimensions {dim}")
 
             unprojected.append(reverse_dims(idct))