torchzero 0.1.8__py3-none-any.whl → 0.3.2__py3-none-any.whl

torchzero/modules/clipping/clipping.py

@@ -0,0 +1,320 @@
+from operator import itemgetter
+from typing import Literal
+from collections.abc import Iterable, Sequence
+import math
+import torch
+
+from ...core import Module, Target, Transform
+from ...utils import NumberList, TensorList, generic_eq
+
+
+def clip_grad_value_(params: Iterable[torch.Tensor], value: float):
+    """Clips gradient of an iterable of parameters at specified value.
+    Gradients are modified in-place.
+    Args:
+        params (Iterable[Tensor]): iterable of tensors with gradients to clip.
+        value (float or int): maximum allowed value of gradient
+    """
+    grads = [p.grad for p in params if p.grad is not None]
+    torch._foreach_clamp_min_(grads, -value)
+    torch._foreach_clamp_max_(grads, value)
+
+def _clip_norm_(
+    tensors_: TensorList,
+    min: float | NumberList | None,
+    max: float | NumberList | None,
+    norm_value: float | NumberList | None,
+    ord: float,
+    dim: int | Sequence[int] | Literal["global"] | None,
+    inverse_dims: bool,
+    min_size: int,
+) -> TensorList:
+    """generic function that can clip norm or normalize"""
+    if norm_value is not None:
+        if min is not None or max is not None:
+            raise ValueError(f'if norm_value is given then min and max must be None got {min = }; {max = }')
+
+        # if dim is None: return tensors_.mul_(norm_value / tensors_.norm(ord=ord))
+        if dim == 'global': return tensors_.mul_(norm_value / tensors_.global_vector_norm(ord=ord))
+
+    # if dim is None: return tensors_.clip_norm_(min,max,tensorwise=True,ord=ord)
+    if dim == 'global': return tensors_.clip_norm_(min,max,tensorwise=False,ord=ord)
+
+    muls = []
+    tensors_to_mul = []
+    if isinstance(dim, int): dim = (dim, )
+
+    for i, tensor in enumerate(tensors_):
+        # remove dimensions that overflow tensor.ndim or are too small
+        if tensor.ndim == 0: tensor = tensor.unsqueeze(0)
+        if dim is None: dim = list(range(tensor.ndim))
+        real_dim = [d for d in dim if d < tensor.ndim]
+        if inverse_dims: real_dim = [d for d in range(tensor.ndim) if d not in real_dim]
+        if len(real_dim) == 0: continue
+        size = math.prod(tensor.size(d) for d in real_dim)
+        if size < min_size: continue
+
+        norm: torch.Tensor = torch.linalg.vector_norm(tensor, ord=ord, dim=real_dim, keepdim=True) # pylint:disable=not-callable
+        if norm.numel() == 1 and norm == 0: continue
+        norm = torch.where(norm == 0, 1, norm)
+
+        # normalize = True, perform normalization
+        norm_v = norm_value[i] if isinstance(norm_value, (list,tuple)) else norm_value
+        if norm_v is not None:
+            mul = norm_v / norm
+
+        # else clip to min and max norms
+        else:
+            minv = min[i] if isinstance(min, (list,tuple)) else min
+            maxv = max[i] if isinstance(max, (list,tuple)) else max
+
+            mul = 1
+            if minv is not None:
+                mul_to_min = (minv / norm).clamp(min=1)
+                mul *= mul_to_min
+
+            if maxv is not None:
+                mul_to_max = (maxv / norm).clamp(max=1)
+                mul *= mul_to_max
+
+        muls.append(mul)
+        tensors_to_mul.append(tensor)
+
+    if len(muls) > 0:
+
+
+        torch._foreach_mul_(tensors_to_mul, muls)
+    return tensors_
+
+
+def clip_grad_norm_(
+    params: Iterable[torch.Tensor],
+    max_norm: float | None,
+    ord: float = 2,
+    dim: int | Sequence[int] | Literal["global"] | None = None,
+    inverse_dims: bool = False,
+    min_size: int = 2,
+    min_norm: float | None = None,
+):
+    """Clips gradient of an iterable of parameters to specified norm value.
+    Gradients are modified in-place.
+
+    Args:
+        params (Iterable[torch.Tensor]): parameters with gradients to clip.
+        value (float): value to clip norm to.
+        ord (float, optional): norm order. Defaults to 2.
+        dim (int | Sequence[int] | str | None, optional):
+            calculates norm along those dimensions.
+            If list/tuple, tensors are normalized along all dimensios in `dim` that they have.
+            Can be set to "global" to normalize by global norm of all gradients concatenated to a vector.
+            Defaults to None.
+        min_size (int, optional):
+            minimal size of a dimension to normalize along it. Defaults to 1.
+    """
+    grads = TensorList(p.grad for p in params if p.grad is not None)
+    _clip_norm_(grads, min=min_norm, max=max_norm, norm_value=None, ord=ord, dim=dim, inverse_dims=inverse_dims, min_size=min_size)
+
+
+def normalize_grads_(
+    params: Iterable[torch.Tensor],
+    norm_value: float,
+    ord: float = 2,
+    dim: int | Sequence[int] | Literal["global"] | None = None,
+    inverse_dims: bool = False,
+    min_size: int = 1,
+):
+    """Normalizes gradient of an iterable of parameters to specified norm value.
+    Gradients are modified in-place.
+
+    Args:
+        params (Iterable[torch.Tensor]): parameters with gradients to clip.
+        norm_value (float): value to clip norm to.
+        ord (float, optional): norm order. Defaults to 2.
+        dim (int | Sequence[int] | str | None, optional):
+            calculates norm along those dimensions.
+            If list/tuple, tensors are normalized along all dimensios in `dim` that they have.
+            Can be set to "global" to normalize by global norm of all gradients concatenated to a vector.
+            Defaults to None.
+        inverse_dims (bool, optional):
+            if True, the `dims` argument is inverted, and all other dimensions are normalized.
+        min_size (int, optional):
+            minimal size of a dimension to normalize along it. Defaults to 1.
+    """
+    grads = TensorList(p.grad for p in params if p.grad is not None)
+    _clip_norm_(grads, min=None, max=None, norm_value=norm_value, ord=ord, dim=dim, inverse_dims=inverse_dims, min_size=min_size)
+
+
+class ClipValue(Transform):
+    """Clips update magnitude to be within `(-value, value)` range."""
+    def __init__(self, value: float, target: Target = 'update'):
+        defaults = dict(value=value)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        value = self.get_settings('value', params=params)
+        return TensorList(tensors).clip_([-v for v in value], value)
+
+class ClipNorm(Transform):
+    """Clips update norm to be no larger than `value`.
+
+    Args:
+        value (float): value to clip norm to.
+        ord (float, optional): norm order. Defaults to 2.
+        dim (int | Sequence[int] | str | None, optional):
+            calculates norm along those dimensions.
+            If list/tuple, tensors are normalized along all dimensios in `dim` that they have.
+            Can be set to "global" to normalize by global norm of all gradients concatenated to a vector.
+            Defaults to None.
+        inverse_dims (bool, optional):
+            if True, the `dims` argument is inverted, and all other dimensions are normalized.
+        min_size (int, optional):
+            minimal numer of elements in a parameter or slice to clip norm. Defaults to 1.
+        target (str, optional):
+            what this affects.
+    """
+    def __init__(
+        self,
+        max_norm: float,
+        ord: float = 2,
+        dim: int | Sequence[int] | Literal["global"] | None = None,
+        inverse_dims: bool = False,
+        min_size: int = 1,
+        target: Target = "update",
+    ):
+        defaults = dict(max_norm=max_norm,ord=ord,dim=dim,min_size=min_size,inverse_dims=inverse_dims)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        max_norm = self.get_settings('max_norm', params=params, cls=NumberList)
+        ord, dim, min_size, inverse_dims = itemgetter('ord', 'dim', 'min_size', 'inverse_dims')(self.settings[params[0]])
+        _clip_norm_(
+            tensors_ = TensorList(tensors),
+            min = 0,
+            max = max_norm,
+            norm_value = None,
+            ord = ord,
+            dim = dim,
+            inverse_dims=inverse_dims,
+            min_size = min_size,
+        )
+        return tensors
+
+class Normalize(Transform):
+    """Normalizes the update.
+
+    Args:
+        value (float): desired norm value.
+        ord (float, optional): norm order. Defaults to 2.
+        dim (int | Sequence[int] | str | None, optional):
+            calculates norm along those dimensions.
+            If list/tuple, tensors are normalized along all dimensios in `dim` that they have.
+            Can be set to "global" to normalize by global norm of all gradients concatenated to a vector.
+            Defaults to None.
+        inverse_dims (bool, optional):
+            if True, the `dims` argument is inverted, and all other dimensions are normalized.
+        min_size (int, optional):
+            minimal size of a dimension to normalize along it. Defaults to 1.
+        target (str, optional):
+            what this affects.
+    """
+    def __init__(
+        self,
+        norm_value: float = 1,
+        ord: float = 2,
+        dim: int | Sequence[int] | Literal["global"] | None = None,
+        inverse_dims: bool = False,
+        min_size: int = 1,
+        target: Target = "update",
+    ):
+        defaults = dict(norm_value=norm_value,ord=ord,dim=dim,min_size=min_size, inverse_dims=inverse_dims)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        norm_value = self.get_settings('norm_value', params=params, cls=NumberList)
+        ord, dim, min_size, inverse_dims = itemgetter('ord', 'dim', 'min_size', 'inverse_dims')(self.settings[params[0]])
+
+        _clip_norm_(
+            tensors_ = TensorList(tensors),
+            min = None,
+            max = None,
+            norm_value = norm_value,
+            ord = ord,
+            dim = dim,
+            inverse_dims=inverse_dims,
+            min_size = min_size,
+        )
+
+        return tensors
+
+
+def _centralize_(
+    tensors_: TensorList,
+    dim: int | Sequence[int] | Literal["global"] | None,
+    min_size: int,
+    inverse_dims: bool,
+) -> TensorList:
+    """generic function that can clip norm or normalize"""
+    if dim == 'global': return tensors_.sub_(tensors_.global_mean().item())
+
+    subs = []
+    tensors_to_sub = []
+    if isinstance(dim, int): dim = (dim, )
+
+    for tensor in tensors_:
+        # remove dimensions that overflow tensor.ndim or are too small
+        if dim is None: dim = list(range(tensor.ndim))
+        real_dim = [d for d in dim if d < tensor.ndim]
+        if inverse_dims: real_dim = [d for d in range(tensor.ndim) if d not in real_dim]
+        if len(real_dim) == 0: continue
+        size = math.prod(tensor.size(d) for d in real_dim)
+        if size < min_size: continue
+
+        mean: torch.Tensor = torch.mean(tensor, dim=real_dim, keepdim=True)
+        if mean.numel() == 1 and mean == 0: continue
+
+        subs.append(mean)
+        tensors_to_sub.append(tensor)
+
+    if len(subs) > 0:
+        torch._foreach_sub_(tensors_to_sub, subs)
+
+    return tensors_
+
+
+class Centralize(Transform):
+    """Centralizes the update.
+
+    Args:
+        value (float): desired norm value.
+        ord (float, optional): norm order. Defaults to 2.
+        dim (int | Sequence[int] | str | None, optional):
+            calculates norm along those dimensions.
+            If list/tuple, tensors are centralized along all dimensios in `dim` that they have.
+            Can be set to "global" to centralize by global mean of all gradients concatenated to a vector.
+            Defaults to None.
+        inverse_dims (bool, optional):
+            if True, the `dims` argument is inverted, and all other dimensions are centralized.
+        min_size (int, optional):
+            minimal size of a dimension to normalize along it. Defaults to 1.
+    """
+    def __init__(
+        self,
+        dim: int | Sequence[int] | Literal["global"] | None = None,
+        inverse_dims: bool = False,
+        min_size: int = 2,
+        target: Target = "update",
+    ):
+        defaults = dict(dim=dim,min_size=min_size,inverse_dims=inverse_dims)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        dim, min_size, inverse_dims = itemgetter('dim', 'min_size', 'inverse_dims')(self.settings[params[0]])
+
+        _centralize_(tensors_ = TensorList(tensors), dim=dim, inverse_dims=inverse_dims, min_size=min_size)
+
+        return tensors
+

torchzero/modules/clipping/ema_clipping.py

@@ -0,0 +1,135 @@
+from operator import itemgetter
+from typing import Literal
+from collections.abc import Iterable, Sequence
+
+import torch
+
+from ...core import Module, Target, Transform, apply, Chainable
+from ...utils import NumberList, TensorList, generic_eq
+
+class ClipNormByEMA(Transform):
+    """Clips norm to be no larger than the norm of an exponential moving average of past updates.
+
+    Args:
+        beta (float, optional): beta for the exponential moving average. Defaults to 0.99.
+        ord (float, optional): order of the norm. Defaults to 2.
+        eps (float, optional): epsilon for division. Defaults to 1e-6.
+        tensorwise (bool, optional): whether to calculate norm separately for each layer, or global norm for all layers. Defaults to True.
+        max_ema_growth (float | None, optional):
+            if specified, exponential moving average norm can grow but at most this value per step. Defaults to 1.5.
+        ema_init (str, optional):
+            How to initialize exponential moving average on first step, "update" to use the first update or "zeros". Defaults to 'zeros'.
+    """
+    NORMALIZE = False
+    def __init__(
+        self,
+        beta=0.99,
+        ord: float = 2,
+        eps=1e-6,
+        tensorwise:bool=True,
+        max_ema_growth: float | None = 1.5,
+        ema_init: Literal['zeros', 'update'] = 'zeros',
+    ):
+        defaults = dict(beta=beta, ord=ord, tensorwise=tensorwise, ema_init=ema_init, eps=eps, max_ema_growth=max_ema_growth)
+        super().__init__(defaults, uses_grad=False)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        ord, tensorwise, ema_init, max_ema_growth = itemgetter('ord', 'tensorwise', 'ema_init', 'max_ema_growth')(self.settings[params[0]])
+
+        beta, eps = self.get_settings('beta', 'eps', params=params, cls=NumberList)
+        tensors = TensorList(tensors)
+
+        ema = self.get_state('ema', params=params, init = (torch.zeros_like if ema_init=='zeros' else tensors), cls=TensorList)
+        ema.lerp_(tensors, 1-beta)
+
+        if tensorwise:
+            ema_norm = ema.norm(ord)
+
+            # clip ema norm growth
+            if max_ema_growth is not None:
+                prev_ema_norm = self.get_state('prev_ema_norm', params=params, init=ema_norm, cls=TensorList)
+                allowed_norm = (prev_ema_norm * max_ema_growth).clip(min=1e-6)
+                ema_denom = (ema_norm / allowed_norm).clip(min=1)
+                ema.div_(ema_denom)
+                ema_norm.div_(ema_denom)
+                prev_ema_norm.set_(ema_norm)
+
+            tensors_norm = tensors.norm(ord)
+            denom = tensors_norm / ema_norm.clip(min=eps)
+            if self.NORMALIZE: denom.clip_(min=eps)
+            else: denom.clip_(min=1)
+
+        else:
+            ema_norm = ema.global_vector_norm(ord)
+
+            # clip ema norm growth
+            if max_ema_growth is not None:
+                prev_ema_norm = self.global_state.setdefault('prev_ema_norm', ema_norm)
+                allowed_norm = prev_ema_norm * max_ema_growth
+                if ema_norm > allowed_norm:
+                    ema.div_(ema_norm / allowed_norm)
+                    ema_norm = allowed_norm
+                prev_ema_norm.set_(ema_norm)
+
+            tensors_norm = tensors.global_vector_norm(ord)
+            denom = tensors_norm / ema_norm.clip(min=eps[0])
+            if self.NORMALIZE: denom.clip_(min=eps[0])
+            else: denom.clip_(min=1)
+
+        tensors.div_(denom)
+        return tensors
+
+class NormalizeByEMA(ClipNormByEMA):
+    """Sets norm of the update to be the same as the norm of an exponential moving average of past updates.
+
+    Args:
+        beta (float, optional): beta for the exponential moving average. Defaults to 0.99.
+        ord (float, optional): order of the norm. Defaults to 2.
+        eps (float, optional): epsilon for division. Defaults to 1e-6.
+        tensorwise (bool, optional): whether to calculate norm separately for each layer, or global norm for all layers. Defaults to True.
+        max_ema_growth (float | None, optional):
+            if specified, exponential moving average norm can grow but at most this value per step. Defaults to 1.5.
+        ema_init (str, optional):
+            How to initialize exponential moving average on first step, "update" to use the first update or "zeros". Defaults to 'zeros'.
+    """
+    NORMALIZE = True
+
+# TODO Centralize by EMA?
+
+class ClipValueByEMA(Transform):
+    """Clips magnitude of update to be no larger than magnitude of an exponential moving average of past (unclipped) updates.
+
+    Args:
+        beta (float, optional): beta for the exponential moving average. Defaults to 0.99.
+        ema_init (str, optional):
+            How to initialize exponential moving average on first step, "update" to use the first update or "zeros". Defaults to 'zeros'.
+        ema_tfm (Chainable | None, optional): optional modules applied to exponential moving average before clipping by it. Defaults to None.
+    """
+    def __init__(
+        self,
+        beta=0.99,
+        ema_init: Literal['zeros', 'update'] = 'zeros',
+        ema_tfm:Chainable | None=None,
+    ):
+        defaults = dict(beta=beta, ema_init=ema_init)
+        super().__init__(defaults, uses_grad=False)
+
+        if ema_tfm is not None:
+            self.set_child('ema_tfm', ema_tfm)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        ema_init = itemgetter('ema_init')(self.settings[params[0]])
+
+        beta = self.get_settings('beta', params=params, cls=NumberList)
+        tensors = TensorList(tensors)
+
+        ema = self.get_state('ema', params=params, init = (torch.zeros_like if ema_init=='zeros' else lambda t: t.abs()), cls=TensorList)
+        ema.lerp_(tensors.abs(), 1-beta)
+
+        if 'ema_tfm' in self.children:
+            ema = TensorList(apply(self.children['ema_tfm'], ema, params, vars.grad, vars))
+
+        tensors.clip_(-ema, ema)
+        return tensors

torchzero/modules/clipping/growth_clipping.py

@@ -0,0 +1,187 @@
+from operator import itemgetter
+
+import torch
+
+from ...core import TensorwiseTransform, Target, Transform
+from ...utils import TensorList, as_tensorlist
+
+
+class ClipValueGrowth(TensorwiseTransform):
+    """Clips update value magnitude growth.
+
+    Args:
+        add (float | None, optional): additive clipping, next update is at most `previous update + add`. Defaults to None.
+        mul (float | None, optional): multiplicative clipping, next update is at most `previous update * mul`. Defaults to 1.5.
+        min_value (float | None, optional):
+            minimum value for multiplicative clipping to prevent collapse to 0.
+            Next update is at most :code:`max(prev_update, min_value) * mul`. Defaults to 1e-4.
+        max_decay (float | None, optional):
+            bounds the tracked multiplicative clipping decay to prevent collapse to 0.
+            Next update is at most :code:`max(previous update * mul, max_decay)`.
+            Defaults to 2.
+        target (Target, optional): what to set on vars.. Defaults to "update".
+    """
+    def __init__(
+        self,
+        add: float | None = None,
+        mul: float | None = 1.5,
+        min_value: float | None = 1e-4,
+        max_decay: float | None = 2,
+        target: Target = "update",
+    ):
+        defaults = dict(add=add, mul=mul, min_value=min_value, max_decay=max_decay)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+
+    def transform(self, tensor, param, grad, vars):
+        add, mul, min_value, max_decay = itemgetter('add','mul','min_value','max_decay')(self.settings[param])
+        add: float | None
+
+        state = self.state[param]
+
+        if add is None and mul is None:
+            return tensor
+
+        if 'prev' not in state:
+            state['prev'] = tensor.clone()
+            return tensor
+
+        prev: torch.Tensor = state['prev']
+
+        # additive bound
+        if add is not None:
+            growth = (tensor.abs() - prev.abs()).clip(min=0)
+            tensor.sub_(torch.where(growth > add, (growth-add).copysign_(tensor), 0))
+
+        # multiplicative bound
+        growth = None
+        if mul is not None:
+            prev_magn = prev.abs()
+            if min_value is not None: prev_magn.clip_(min=min_value)
+            growth = (tensor.abs() / prev_magn).clamp_(min=1e-8)
+
+            denom = torch.where(growth > mul, growth/mul, 1)
+
+            tensor.div_(denom)
+
+        # limit max growth decay
+        if max_decay is not None:
+            if growth is None:
+                prev_magn = prev.abs()
+                if min_value is not None: prev_magn.clip_(min=min_value)
+                growth = (tensor.abs() / prev_magn).clamp_(min=1e-8)
+
+            new_prev = torch.where(growth < (1/max_decay), prev/max_decay, tensor)
+        else:
+            new_prev = tensor.clone()
+
+        state['prev'] = new_prev
+        return tensor
+
+
+def norm_growth_clip_(
+    tensor_: torch.Tensor,
+    prev_norm: torch.Tensor,
+    add: float | None,
+    mul: float | None,
+    min_value: float | None,
+    max_decay: float | None,
+    ord: float,
+):
+    if add is None and mul is None: return tensor_
+    norm = torch.linalg.vector_norm(tensor_, ord=ord) # pylint:disable=not-callable
+
+    denom = 1
+    # additive bound
+    if add is not None:
+        allowed_norm = prev_norm + add
+        if norm > allowed_norm: denom = norm / allowed_norm
+
+    # multiplicative bound
+    if mul is not None:
+        allowed_norm = prev_norm * mul
+        if norm > allowed_norm: denom = max(denom, norm / allowed_norm)
+
+    # minimal norm
+    if min_value is not None:
+        denom = max(denom, min_value)
+
+    # limit max growth decay
+    new_prev_norm = norm/denom
+    if max_decay is not None:
+        decay = norm / prev_norm
+        if decay < (1/max_decay):
+            new_prev_norm = prev_norm / max_decay
+
+    if min_value is not None: new_prev_norm = max(new_prev_norm, min_value) # pyright:ignore[reportArgumentType]
+    return tensor_.div_(denom), new_prev_norm, denom
+
+
+class ClipNormGrowth(Transform):
+    """Clips update norm growth.
+
+    Args:
+        add (float | None, optional): additive clipping, next update norm is at most `previous norm + add`. Defaults to None.
+        mul (float | None, optional): multiplicative clipping, next update norm is at most `previous norm * mul`. Defaults to 1.5.
+        min_value (float | None, optional):
+            minimum value for multiplicative clipping to prevent collapse to 0.
+            Next norm is at most :code:`max(prev_norm, min_value) * mul`. Defaults to 1e-4.
+        max_decay (float | None, optional):
+            bounds the tracked multiplicative clipping decay to prevent collapse to 0.
+            Next norm is at most :code:`max(previous norm * mul, max_decay)`.
+            Defaults to 2.
+        ord (float, optional): norm order. Defaults to 2.
+        parameterwise (bool, optional):
+            if True, norms are calculated parameter-wise, otherwise treats all parameters as single vector. Defaults to True.
+        target (Target, optional): what to set on vars. Defaults to "update".
+    """
+    def __init__(
+        self,
+        add: float | None = None,
+        mul: float | None = 1.5,
+        min_value: float | None = 1e-4,
+        max_decay: float | None = 2,
+        ord: float = 2,
+        parameterwise=True,
+        target: Target = "update",
+    ):
+        defaults = dict(add=add, mul=mul, min_value=min_value, max_decay=max_decay, ord=ord, parameterwise=parameterwise)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+
+
+    def transform(self, tensors, params, grads, vars):
+        parameterwise = self.settings[params[0]]['parameterwise']
+        tensors = TensorList(tensors)
+
+        if parameterwise:
+            ts = tensors
+            stts = [self.state[p] for p in params]
+            stns = [self.settings[p] for p in params]
+
+        else:
+            ts = [tensors.to_vec()]
+            stts = [self.global_state]
+            stns = [self.settings[params[0]]]
+
+
+        for t,state, settings in zip(ts, stts, stns):
+            if 'prev_norm' not in state:
+                state['prev_norm'] = torch.linalg.vector_norm(t, ord=settings['ord']) # pylint:disable=not-callable
+                state['prev_denom'] = 1
+                continue
+
+            _,  state['prev_norm'], state['prev_denom'] = norm_growth_clip_(
+                tensor_ = t,
+                prev_norm = state['prev_norm'],
+                add = settings['add'],
+                mul = settings['mul'],
+                min_value = settings['min_value'],
+                max_decay = settings['max_decay'],
+                ord = settings['ord'],
+            )
+
+        if not parameterwise:
+            tensors.from_vec_(ts[0])
+
+        return tensors

torchzero/modules/experimental/__init__.py

@@ -1,19 +1,14 @@
-"""Optimizers that I haven't tested and various (mostly stupid) ideas go there.
-If something works well I will move it outside of experimental folder.
-Otherwise all optimizers in this category should be considered unlikely to good for most tasks."""
-from .experimental import GradMin, HVPDiagNewton, MinibatchRprop, ReduceOutwardLR
-from .quad_interp import QuadraticInterpolation2Point
-from .subspace import (
-    Proj2Masks,
-    ProjAscent,
-    ProjAscentRay,
-    Projection,
-    ProjGrad,
-    ProjGradAscentDifference,
-    ProjGradRay,
-    ProjLastAscentDifference,
-    ProjLastGradDifference,
-    ProjNormalize,
-    ProjRandom,
-    Subspace,
+from .absoap import ABSOAP
+from .adadam import Adadam
+from .adamY import AdamY
+from .adasoap import AdaSOAP
+from .curveball import CurveBall
+from .dsoap import DSOAP
+from .gradmin import GradMin
+from .reduce_outward_lr import ReduceOutwardLR
+from .spectral import SpectralPreconditioner
+from .subspace_preconditioners import (
+    HistorySubspacePreconditioning,
+    RandomSubspacePreconditioning,
 )
+from .tropical_newton import TropicalNewton