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

torchzero/modules/experimental/spectral.py

@@ -1,288 +0,0 @@
-from abc import ABC, abstractmethod
-import math
-from collections import deque
-from typing import Literal, Any
-
-import torch
-from ...core import Chainable, TensorwisePreconditioner
-from ...utils.linalg.matrix_funcs import matrix_power_eigh
-from ...utils.linalg.svd import randomized_svd
-from ...utils.linalg.qr import qr_householder
-
-
-class _Solver:
-    @abstractmethod
-    def update(self, history: deque[torch.Tensor], damping: float | None) -> tuple[Any, Any]:
-        """returns stuff for apply"""
-    @abstractmethod
-    def apply(self, __g: torch.Tensor, __A:torch.Tensor, __B:torch.Tensor) -> torch.Tensor:
-        """apply preconditioning to tensor"""
-
-class _SVDSolver(_Solver):
-    def __init__(self, driver=None): self.driver=driver
-    def update(self, history, damping):
-        M_hist = torch.stack(tuple(history), dim=1)
-        device = None # driver is CUDA only
-        if self.driver is not None:
-            device = M_hist.device
-            M_hist = M_hist.cuda()
-
-        try:
-            U, S, _ = torch.linalg.svd(M_hist, full_matrices=False, driver=self.driver) # pylint:disable=not-callable
-
-            if self.driver is not None:
-                U = U.to(device); S = S.to(device)
-
-            if damping is not None and damping != 0: S.add_(damping)
-            return U, S
-
-        except torch.linalg.LinAlgError:
-            return None, None
-
-    def apply(self, g: torch.Tensor, U: torch.Tensor, S: torch.Tensor):
-        Utg = (U.T @ g).div_(S)
-        return U @ Utg
-
-class _SVDLowRankSolver(_Solver):
-    def __init__(self, q: int = 6, niter: int = 2): self.q, self.niter = q, niter
-    def update(self, history, damping):
-        M_hist = torch.stack(tuple(history), dim=1)
-        try:
-            U, S, _ = torch.svd_lowrank(M_hist, q=self.q, niter=self.niter)
-            if damping is not None and damping != 0: S.add_(damping)
-            return U, S
-        except torch.linalg.LinAlgError:
-            return None, None
-
-    def apply(self, g: torch.Tensor, U: torch.Tensor, S: torch.Tensor):
-        Utg = (U.T @ g).div_(S)
-        return U @ Utg
-
-class _RandomizedSVDSolver(_Solver):
-    def __init__(self, k: int = 3, driver: str | None = 'gesvda'):
-        self.driver = driver
-        self.k = k
-
-    def update(self, history, damping):
-        M_hist = torch.stack(tuple(history), dim=1)
-        device = None # driver is CUDA only
-        if self.driver is not None:
-            device = M_hist.device
-            M_hist = M_hist.cuda()
-
-        try:
-            U, S, _ = randomized_svd(M_hist, k=self.k, driver=self.driver)
-
-            if self.driver is not None:
-                U = U.to(device); S = S.to(device)
-
-            if damping is not None and damping != 0: S.add_(damping)
-            return U, S
-
-        except torch.linalg.LinAlgError:
-            return None, None
-
-    def apply(self, g: torch.Tensor, U: torch.Tensor, S: torch.Tensor):
-        Utg = (U.T @ g).div_(S)
-        return U @ Utg
-
-class _QRDiagonalSolver(_Solver):
-    def __init__(self, sqrt=True): self.sqrt = sqrt
-    def update(self, history, damping):
-        M_hist = torch.stack(tuple(history), dim=1)
-        try:
-            Q, R = torch.linalg.qr(M_hist, mode='reduced') # pylint:disable=not-callable
-            R_diag = R.diag().abs()
-            if damping is not None and damping != 0: R_diag.add_(damping)
-            if self.sqrt: R_diag.sqrt_()
-            return Q, R_diag
-        except torch.linalg.LinAlgError:
-            return None, None
-
-    def apply(self, g: torch.Tensor, Q: torch.Tensor, R_diag: torch.Tensor):
-        Qtg = (Q.T @ g).div_(R_diag)
-        return Q @ Qtg
-
-class _QRSolver(_Solver):
-    def __init__(self, sqrt=True): self.sqrt = sqrt
-    def update(self, history, damping):
-        M_hist = torch.stack(tuple(history), dim=1)
-        try:
-            # Q: d x k, R: k x k
-            Q, R = torch.linalg.qr(M_hist, mode='reduced') # pylint:disable=not-callable
-            A = R @ R.T
-            if damping is not None and damping != 0: A.diagonal(dim1=-2, dim2=-1).add_(damping)
-            if self.sqrt: A = matrix_power_eigh(A, 0.5)
-            return Q, A
-        except (torch.linalg.LinAlgError):
-            return None,None
-
-    def apply(self, g: torch.Tensor, Q: torch.Tensor, A: torch.Tensor) -> torch.Tensor:
-        g_proj = Q.T @ g
-        y, _ = torch.linalg.solve_ex(A, g_proj) # pylint:disable=not-callable
-        return Q @ y
-
-class _QRHouseholderSolver(_Solver):
-    def __init__(self, sqrt=True): self.sqrt = sqrt
-    def update(self, history, damping):
-        M_hist = torch.stack(tuple(history), dim=1)
-        try:
-            # Q: d x k, R: k x k
-            Q, R = qr_householder(M_hist, mode='reduced') # pylint:disable=not-callable
-            A = R @ R.T
-            if damping is not None and damping != 0: A.diagonal(dim1=-2, dim2=-1).add_(damping)
-            if self.sqrt: A = matrix_power_eigh(A, 0.5)
-            return Q, A
-        except (torch.linalg.LinAlgError):
-            return None,None
-
-    def apply(self, g: torch.Tensor, Q: torch.Tensor, A: torch.Tensor) -> torch.Tensor:
-        g_proj = Q.T @ g
-        y, _ = torch.linalg.solve_ex(A, g_proj) # pylint:disable=not-callable
-        return Q @ y
-
-
-class _EighSolver(_Solver):
-    def __init__(self, sqrt=True):
-        self.sqrt = sqrt
-
-    def update(self, history, damping):
-        M_hist = torch.stack(tuple(history), dim=1)
-        grams = M_hist @ M_hist.T # (d, d)
-        if damping is not None and damping != 0: grams.diagonal(dim1=-2, dim2=-1).add_(damping)
-        try:
-            L, Q = torch.linalg.eigh(grams) # L: (d,), Q: (d, d) # pylint:disable=not-callable
-            L = L.abs().clamp_(min=1e-12)
-            if self.sqrt: L = L.sqrt()
-            return Q, L
-        except torch.linalg.LinAlgError:
-            return None, None
-
-    def apply(self, g: torch.Tensor, Q: torch.Tensor, L: torch.Tensor) -> torch.Tensor:
-        Qtg = (Q.T @ g).div_(L)
-        return Q @ Qtg
-
-
-SOLVERS = {
-    "svd": _SVDSolver(), # fallbacks on "gesvd" which basically takes ages or just hangs completely
-    "svd_gesvdj": _SVDSolver("gesvdj"), # no fallback on slow "gesvd"
-    "svd_gesvda": _SVDSolver("gesvda"), # approximate method for wide matrices, sometimes better sometimes worse but faster
-    "svd_lowrank": _SVDLowRankSolver(), # maybe need to tune parameters for this, with current ones its slower and worse
-    "randomized_svd2": _RandomizedSVDSolver(2),
-    "randomized_svd3": _RandomizedSVDSolver(3),
-    "randomized_svd4": _RandomizedSVDSolver(4),
-    "randomized_svd5": _RandomizedSVDSolver(5),
-    "eigh": _EighSolver(), # this is O(n**2) storage, but is this more accurate?
-    "qr": _QRSolver(),
-    "qr_householder": _QRHouseholderSolver(), # this is slower... but maybe it won't freeze? I think svd_gesvda is better
-    "qrdiag": _QRDiagonalSolver(),
-}
-
-def maybe_lerp_(state_: dict, beta: float | None, key, value: Any):
-    if (key not in state_) or (beta is None) or (not isinstance(value, torch.Tensor)): state_[key] = value
-    else:
-        if state_[key].shape != value.shape: state_[key] = value
-        else: state_[key].lerp_(value, 1-beta)
-
-class SpectralPreconditioner(TensorwisePreconditioner):
-    """Whitening preconditioner via SVD on history of past gradients or gradient differences scaled by parameter differences.
-
-    Args:
-        history_size (int, optional): number of past gradients to store for preconditioning. Defaults to 10.
-        update_freq (int, optional): how often to re-compute the preconditioner. Defaults to 1.
-        damping (float, optional): damping term, makes it closer to GD. Defaults to 1e-7.
-        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'.
-        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.
-        scale_first (bool, optional): makes first step small, usually not needed. Defaults to False.
-        inner (Chainable | None, optional): Inner modules applied after updating preconditioner and before applying it. Defaults to None.
-    """
-    def __init__(
-        self,
-        history_size: int = 10,
-        update_freq: int = 1,
-        damping: float = 1e-12,
-        order: int = 1,
-        solver: Literal['svd', 'svd_gesvdj', 'svd_gesvda', 'svd_lowrank', 'eigh', 'qr', 'qrdiag', 'qr_householder'] | _Solver | str = 'svd_gesvda',
-        A_beta: float | None = None,
-        B_beta: float | None = None,
-        interval: int = 1,
-        concat_params: bool = False,
-        scale_first: bool = False,
-        inner: Chainable | None = None,
-    ):
-        if isinstance(solver, str): solver = SOLVERS[solver]
-        # history is still updated each step so Precondition's update_freq has different meaning
-        defaults = dict(history_size=history_size, update_freq=update_freq, damping=damping, order=order, A_beta=A_beta, B_beta=B_beta, solver=solver)
-        super().__init__(defaults, uses_grad=False, concat_params=concat_params, scale_first=scale_first, inner=inner, update_freq=interval)
-
-    @torch.no_grad
-    def update_tensor(self, tensor, param, grad, state, settings):
-        order = settings['order']
-        history_size = settings['history_size']
-        update_freq = settings['update_freq']
-        damping = settings['damping']
-        A_beta = settings['A_beta']
-        B_beta = settings['B_beta']
-        solver: _Solver = settings['solver']
-
-        if 'history' not in state: state['history'] = deque(maxlen=history_size)
-        history = state['history']
-
-        if order == 1: history.append(tensor.clone().view(-1))
-        else:
-
-            # if order=2, history is of gradient differences, order 3 is differences between differences, etc
-            # normalized by parameter differences
-            cur_p = param.clone()
-            cur_g = tensor.clone()
-            for i in range(1, order):
-                if f'prev_g_{i}' not in state:
-                    state[f'prev_p_{i}'] = cur_p
-                    state[f'prev_g_{i}'] = cur_g
-                    break
-
-                s_k = cur_p - state[f'prev_p_{i}']
-                y_k = cur_g - state[f'prev_g_{i}']
-                state[f'prev_p_{i}'] = cur_p
-                state[f'prev_g_{i}'] = cur_g
-                cur_p = s_k
-                cur_g = y_k
-
-                if i == order - 1:
-                    cur_g = cur_g / torch.linalg.norm(cur_p).clip(min=1e-8) # pylint:disable=not-callable
-                    history.append(cur_g.view(-1))
-
-        step = state.get('step', 0)
-        if step % update_freq == 0 and len(history) != 0:
-            A, B = solver.update(history, damping=damping)
-            maybe_lerp_(state, A_beta, 'A', A)
-            maybe_lerp_(state, B_beta, 'B', B)
-
-        if len(history) != 0:
-            state['step'] = step + 1 # do not increment if no history (gathering s_ks and y_ks)
-
-    @torch.no_grad
-    def apply_tensor(self, tensor, param, grad, state, settings):
-        history_size = settings['history_size']
-        solver: _Solver = settings['solver']
-
-        A = state.get('A', None)
-        if A is None:
-            # make a conservative step to avoid issues due to different GD scaling
-            return tensor.clip_(-0.1, 0.1) # pyright:ignore[reportArgumentType]
-
-        B = state['B']
-        update = solver.apply(tensor.view(-1), A, B).view_as(tensor)
-
-        n = len(state['history'])
-        if n != history_size: update.mul_(n/history_size)
-        return update
-

torchzero/modules/experimental/structured_newton.py

@@ -1,111 +0,0 @@
-# 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/tropical_newton.py

@@ -1,136 +0,0 @@
-import warnings
-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 ...utils.derivatives import (
-    hessian_list_to_mat,
-    hessian_mat,
-    jacobian_and_hessian_wrt,
-)
-from ..second_order.newton import lu_solve, cholesky_solve, least_squares_solve
-
-def tropical_sum(x, dim): return torch.amax(x, dim=dim)
-def tropical_mul(x, y): return x+y
-
-def tropical_matmul(x: torch.Tensor, y: torch.Tensor):
-    # this imlements matmul by calling mul and sum
-
-    x_squeeze = False
-    y_squeeze = False
-
-    if x.ndim == 1:
-        x_squeeze = True
-        x = x.unsqueeze(0)
-
-    if y.ndim == 1:
-        y_squeeze = True
-        y = y.unsqueeze(1)
-
-    res = tropical_sum(tropical_mul(x.unsqueeze(-1), y.unsqueeze(-3)), dim = -2)
-
-    if x_squeeze: res = res.squeeze(-2)
-    if y_squeeze: res = res.squeeze(-1)
-
-    return res
-
-def tropical_dot(x:torch.Tensor, y:torch.Tensor):
-    assert x.ndim == 1 and y.ndim == 1
-    return tropical_matmul(x.unsqueeze(0), y.unsqueeze(1))
-
-def tropical_outer(x:torch.Tensor, y:torch.Tensor):
-    assert x.ndim == 1 and y.ndim == 1
-    return tropical_matmul(x.unsqueeze(1), y.unsqueeze(0))
-
-
-def tropical_solve(A: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
-    r = b.unsqueeze(1) - A
-    return r.amin(dim=-2)
-
-def tropical_solve_and_reconstruct(A: torch.Tensor, b: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-    r = b.unsqueeze(1) - A
-    x = r.amin(dim=-2)
-    b_hat = tropical_matmul(A, x)
-    return x, b_hat
-
-def tikhonov(H: torch.Tensor, reg: float):
-    if reg!=0: H += torch.eye(H.size(-1), dtype=H.dtype, device=H.device) * reg
-    return H
-
-
-class TropicalNewton(Module):
-    """suston"""
-    def __init__(
-        self,
-        reg: float | None = None,
-        hessian_method: Literal["autograd", "func", "autograd.functional"] = "autograd",
-        vectorize: bool = True,
-        interpolate:bool=False,
-        inner: Chainable | None = None,
-    ):
-        defaults = dict(reg=reg, hessian_method=hessian_method, vectorize=vectorize, interpolate=interpolate)
-        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']
-        hessian_method = settings['hessian_method']
-        vectorize = settings['vectorize']
-        interpolate = settings['interpolate']
-
-        # ------------------------ calculate grad and hessian ------------------------ #
-        if hessian_method == 'autograd':
-            with torch.enable_grad():
-                loss = vars.loss = vars.loss_approx = closure(False)
-                g_list, H_list = jacobian_and_hessian_wrt([loss], params, batched=vectorize)
-                g_list = [t[0] for t in g_list] # remove leading dim from loss
-                vars.grad = g_list
-                H = hessian_list_to_mat(H_list)
-
-        elif hessian_method in ('func', 'autograd.functional'):
-            strat = 'forward-mode' if vectorize else 'reverse-mode'
-            with torch.enable_grad():
-                g_list = vars.get_grad(retain_graph=True)
-                H: torch.Tensor = hessian_mat(partial(closure, backward=False), params,
-                                method=hessian_method, vectorize=vectorize, outer_jacobian_strategy=strat) # pyright:ignore[reportAssignmentType]
-
-        else:
-            raise ValueError(hessian_method)
-
-        # -------------------------------- inner step -------------------------------- #
-        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])
-
-        # ------------------------------- regulazition ------------------------------- #
-        if reg is not None: H = tikhonov(H, reg)
-
-        # ----------------------------------- solve ---------------------------------- #
-        tropical_update, g_hat = tropical_solve_and_reconstruct(H, g)
-
-        g_norm = torch.linalg.vector_norm(g) # pylint:disable=not-callable
-        abs_error = torch.linalg.vector_norm(g-g_hat) # pylint:disable=not-callable
-        rel_error = abs_error/g_norm.clip(min=1e-8)
-
-        if interpolate:
-            if rel_error > 1e-8:
-
-                update = cholesky_solve(H, g)
-                if update is None: update = lu_solve(H, g)
-                if update is None: update = least_squares_solve(H, g)
-
-                tropical_update.lerp_(update.ravel(), rel_error.clip(max=1))
-
-        vars.update = vec_to_tensors(tropical_update, params)
-        return vars

torchzero/modules/lr/__init__.py

@@ -1,2 +0,0 @@
-from .lr import LR, StepSize, Warmup
-from .step_size import PolyakStepSize, RandomStepSize

torchzero/modules/lr/lr.py

@@ -1,59 +0,0 @@
-import torch
-
-from ...core import Transform
-from ...utils import NumberList, TensorList, generic_eq
-
-
-def lazy_lr(tensors: TensorList, lr: float | list, inplace:bool):
-    """multiplies by lr if lr is not 1"""
-    if generic_eq(lr, 1): return tensors
-    if inplace: return tensors.mul_(lr)
-    return tensors * lr
-
-class LR(Transform):
-    def __init__(self, lr: float):
-        defaults=dict(lr=lr)
-        super().__init__(defaults, uses_grad=False)
-
-    @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        return lazy_lr(TensorList(tensors), lr=self.get_settings('lr', params=params), inplace=True)
-
-class StepSize(Transform):
-    """this is exactly the same as LR, except the `lr` parameter can be renamed to any other name"""
-    def __init__(self, step_size: float, key = 'step_size'):
-        defaults={"key": key, key: step_size}
-        super().__init__(defaults, uses_grad=False)
-
-    @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        lrs = []
-        for p in params:
-            settings = self.settings[p]
-            lrs.append(settings[settings['key']])
-        return lazy_lr(TensorList(tensors), lr=lrs, inplace=True)
-
-
-def warmup(step: int, start_lr: float | NumberList, end_lr: float | NumberList, steps: float):
-    """returns warm up lr scalar"""
-    if step > steps: return end_lr
-    return start_lr + (end_lr - start_lr) * (step / steps)
-
-class Warmup(Transform):
-    def __init__(self, start_lr = 1e-5, end_lr:float = 1, steps = 100):
-        defaults = dict(start_lr=start_lr,end_lr=end_lr, steps=steps)
-        super().__init__(defaults, uses_grad=False)
-
-    @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        start_lr, end_lr = self.get_settings('start_lr', 'end_lr', params=params, cls = NumberList)
-        num_steps = self.settings[params[0]]['steps']
-        step = self.global_state.get('step', 0)
-
-        target = lazy_lr(
-            TensorList(tensors),
-            lr=warmup(step=step, start_lr=start_lr, end_lr=end_lr, steps=num_steps),
-            inplace=True
-        )
-        self.global_state['step'] = step + 1
-        return target

torchzero/modules/lr/step_size.py

@@ -1,97 +0,0 @@
-import random
-from typing import Any
-
-import torch
-
-from ...core import Transform
-from ...utils import TensorList, NumberList
-
-
-class PolyakStepSize(Transform):
-    """Polyak step-size.
-
-    Args:
-        max (float | None, optional): maximum possible step size. Defaults to None.
-        min_obj_value (int, optional): (estimated) minimal possible value of the objective function (lowest possible loss). Defaults to 0.
-        use_grad (bool, optional):
-            if True, uses dot product of update and gradient to compute the step size.
-            Otherwise, dot product of update with itself is used, which has no geometric meaning so it probably won't work well.
-            Defaults to True.
-        parameterwise (bool, optional):
-            if True, calculate Polyak step-size for each parameter separately,
-            if False calculate one global step size for all parameters. Defaults to False.
-        alpha (float, optional): multiplier to Polyak step-size. Defaults to 1.
-    """
-    def __init__(self, max: float | None = None, min_obj_value: float = 0, use_grad=True, parameterwise=False, alpha: float = 1):
-
-        defaults = dict(alpha=alpha, max=max, min_obj_value=min_obj_value, use_grad=use_grad, parameterwise=parameterwise)
-        super().__init__(defaults, uses_grad=use_grad)
-
-    @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        loss = vars.get_loss(False)
-        assert grads is not None
-        tensors = TensorList(tensors)
-        grads = TensorList(grads)
-        alpha = self.get_settings('alpha', params=params, cls=NumberList)
-        settings = self.settings[params[0]]
-        parameterwise = settings['parameterwise']
-        use_grad = settings['use_grad']
-        max = settings['max']
-        min_obj_value = settings['min_obj_value']
-
-        if parameterwise:
-            if use_grad: denom = (tensors * grads).sum()
-            else: denom = tensors.pow(2).sum()
-            polyak_step_size: TensorList | Any = (loss - min_obj_value) / denom.where(denom!=0, 1)
-            polyak_step_size = polyak_step_size.where(denom != 0, 0)
-            if max is not None: polyak_step_size = polyak_step_size.clamp_max(max)
-
-        else:
-            if use_grad: denom = tensors.dot(grads)
-            else: denom = tensors.dot(tensors)
-            if denom == 0: polyak_step_size = 0 # we converged
-            else: polyak_step_size = (loss - min_obj_value) / denom
-
-            if max is not None:
-                if polyak_step_size > max: polyak_step_size = max
-
-        tensors.mul_(alpha * polyak_step_size)
-        return tensors
-
-
-
-class RandomStepSize(Transform):
-    """Uses random global step size from `low` to `high`.
-
-    Args:
-        low (float, optional): minimum learning rate. Defaults to 0.
-        high (float, optional): maximum learning rate. Defaults to 1.
-        parameterwise (bool, optional):
-            if True, generate random step size for each parameter separately,
-            if False generate one global random step size. Defaults to False.
-    """
-    def __init__(self, low: float = 0, high: float = 1, parameterwise=False, seed:int|None=None):
-        defaults = dict(low=low, high=high, parameterwise=parameterwise,seed=seed)
-        super().__init__(defaults, uses_grad=False)
-
-    @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        settings = self.settings[params[0]]
-        parameterwise = settings['parameterwise']
-
-        seed = settings['seed']
-        if 'generator' not in self.global_state:
-            self.global_state['generator'] = random.Random(seed)
-        generator: random.Random = self.global_state['generator']
-
-        if parameterwise:
-            low, high = self.get_settings('low', 'high', params=params)
-            lr = [generator.uniform(l, h) for l, h in zip(low, high)]
-        else:
-            low = settings['low']
-            high = settings['high']
-            lr = generator.uniform(low, high)
-
-        torch._foreach_mul_(tensors, lr)
-        return tensors