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

torchzero/modules/ops/reduce.py

@@ -5,10 +5,10 @@ from typing import Any, cast
 
 import torch
 
-from ...core import Chainable, Module, Target, Vars, maybe_chain
+from ...core import Chainable, Module, Target, Var, maybe_chain
 
 
-class ReduceOperation(Module, ABC):
+class ReduceOperationBase(Module, ABC):
     """Base class for reduction operations like Sum, Prod, Maximum. This is an abstract class, subclass it and override `transform` method to use it."""
     def __init__(self, defaults: dict[str, Any] | None, *operands: Chainable | Any):
         super().__init__(defaults=defaults)
@@ -26,33 +26,34 @@ class ReduceOperation(Module, ABC):
             raise ValueError('At least one operand must be a module')
 
     @abstractmethod
-    def transform(self, vars: Vars, *operands: Any | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, var: Var, *operands: Any | list[torch.Tensor]) -> list[torch.Tensor]:
         """applies the operation to operands"""
         raise NotImplementedError
 
     @torch.no_grad
-    def step(self, vars: Vars) -> Vars:
+    def step(self, var: Var) -> Var:
         # pass cloned update to all module operands
         processed_operands: list[Any | list[torch.Tensor]] = self.operands.copy()
 
         for i, v in enumerate(self.operands):
             if f'operand_{i}' in self.children:
                 v: Module
-                updated_vars = v.step(vars.clone(clone_update=True))
-                processed_operands[i] = updated_vars.get_update()
-                vars.update_attrs_from_clone_(updated_vars) # update loss, grad, etc if this module calculated them
+                updated_var = v.step(var.clone(clone_update=True))
+                processed_operands[i] = updated_var.get_update()
+                var.update_attrs_from_clone_(updated_var) # update loss, grad, etc if this module calculated them
 
-        transformed = self.transform(vars, *processed_operands)
-        vars.update = transformed
-        return vars
+        transformed = self.transform(var, *processed_operands)
+        var.update = transformed
+        return var
 
-class Sum(ReduceOperation):
+class Sum(ReduceOperationBase):
+    """Outputs sum of :code:`inputs` that can be modules or numbers."""
     USE_MEAN = False
     def __init__(self, *inputs: Chainable | float):
         super().__init__({}, *inputs)
 
     @torch.no_grad
-    def transform(self, vars: Vars, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, var: Var, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
         sorted_inputs = sorted(inputs, key=lambda x: isinstance(x, float))
         sum = cast(list, sorted_inputs[0])
         if len(sorted_inputs) > 1:
@@ -63,12 +64,14 @@ class Sum(ReduceOperation):
         return sum
 
 class Mean(Sum):
+    """Outputs a mean of :code:`inputs` that can be modules or numbers."""
     USE_MEAN = True
 
 
-class WeightedSum(ReduceOperation):
+class WeightedSum(ReduceOperationBase):
     USE_MEAN = False
     def __init__(self, *inputs: Chainable | float, weights: Iterable[float]):
+        """Outputs a weighted sum of :code:`inputs` that can be modules or numbers."""
         weights = list(weights)
         if len(inputs) != len(weights):
             raise ValueError(f'Number of inputs {len(inputs)} must match number of weights {len(weights)}')
@@ -76,9 +79,9 @@ class WeightedSum(ReduceOperation):
         super().__init__(defaults=defaults, *inputs)
 
     @torch.no_grad
-    def transform(self, vars: Vars, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, var: Var, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
         sorted_inputs = sorted(inputs, key=lambda x: isinstance(x, float))
-        weights = self.settings[vars.params[0]]['weights']
+        weights = self.settings[var.params[0]]['weights']
         sum = cast(list, sorted_inputs[0])
         torch._foreach_mul_(sum, weights[0])
         if len(sorted_inputs) > 1:
@@ -91,14 +94,16 @@ class WeightedSum(ReduceOperation):
 
 
 class WeightedMean(WeightedSum):
+    """Outputs weighted mean of :code:`inputs` that can be modules or numbers."""
     USE_MEAN = True
 
-class Median(ReduceOperation):
+class Median(ReduceOperationBase):
+    """Outputs median of :code:`inputs` that can be modules or numbers."""
     def __init__(self, *inputs: Chainable | float):
         super().__init__({}, *inputs)
 
     @torch.no_grad
-    def transform(self, vars: Vars, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, var: Var, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
         res = []
         lists = [i for i in inputs if isinstance(i, list)]
         floats = [i for i in inputs if isinstance(i, (int,float))]
@@ -106,12 +111,13 @@ class Median(ReduceOperation):
             res.append(torch.median(torch.stack(tensors + tuple(torch.full_like(tensors[0], f) for f in floats)), dim=0))
         return res
 
-class Prod(ReduceOperation):
+class Prod(ReduceOperationBase):
+    """Outputs product of :code:`inputs` that can be modules or numbers."""
     def __init__(self, *inputs: Chainable | float):
         super().__init__({}, *inputs)
 
     @torch.no_grad
-    def transform(self, vars: Vars, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, var: Var, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
         sorted_inputs = sorted(inputs, key=lambda x: isinstance(x, float))
         prod = cast(list, sorted_inputs[0])
         if len(sorted_inputs) > 1:
@@ -120,12 +126,13 @@ class Prod(ReduceOperation):
 
         return prod
 
-class MaximumModules(ReduceOperation):
+class MaximumModules(ReduceOperationBase):
+    """Outputs elementwise maximum of :code:`inputs` that can be modules or numbers."""
     def __init__(self, *inputs: Chainable | float):
         super().__init__({}, *inputs)
 
     @torch.no_grad
-    def transform(self, vars: Vars, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, var: Var, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
         sorted_inputs = sorted(inputs, key=lambda x: isinstance(x, float))
         maximum = cast(list, sorted_inputs[0])
         if len(sorted_inputs) > 1:
@@ -134,12 +141,13 @@ class MaximumModules(ReduceOperation):
 
         return maximum
 
-class MinimumModules(ReduceOperation):
+class MinimumModules(ReduceOperationBase):
+    """Outputs elementwise minimum of :code:`inputs` that can be modules or numbers."""
     def __init__(self, *inputs: Chainable | float):
         super().__init__({}, *inputs)
 
     @torch.no_grad
-    def transform(self, vars: Vars, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, var: Var, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
         sorted_inputs = sorted(inputs, key=lambda x: isinstance(x, float))
         minimum = cast(list, sorted_inputs[0])
         if len(sorted_inputs) > 1:

torchzero/modules/ops/unary.py

@@ -3,79 +3,95 @@ from collections import deque
 import torch
 
 from ...core import TensorwiseTransform, Target, Transform
-from ...utils import TensorList
+from ...utils import TensorList, unpack_dicts,unpack_states
 
 class UnaryLambda(Transform):
+    """Applies :code:`fn` to input tensors.
+
+    :code:`fn` must accept and return a list of tensors.
+    """
     def __init__(self, fn, target: "Target" = 'update'):
         defaults = dict(fn=fn)
         super().__init__(defaults=defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        return self.settings[params[0]]['fn'](tensors)
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        return settings[0]['fn'](tensors)
 
 class UnaryParameterwiseLambda(TensorwiseTransform):
+    """Applies :code:`fn` to each input tensor.
+
+    :code:`fn` must accept and return a tensor.
+    """
     def __init__(self, fn, target: "Target" = 'update'):
         defaults = dict(fn=fn)
         super().__init__(uses_grad=False, defaults=defaults, target=target)
 
     @torch.no_grad
-    def transform(self, tensor, param, grad, vars):
-        return self.settings[param]['fn'](tensor)
+    def apply_tensor(self, tensor, param, grad, loss, state, setting):
+        return setting['fn'](tensor)
 
 class CustomUnaryOperation(Transform):
+    """Applies :code:`getattr(tensor, name)` to each tensor
+    """
     def __init__(self, name: str, target: "Target" = 'update'):
         defaults = dict(name=name)
         super().__init__(defaults=defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        return getattr(tensors, self.settings[params[0]]['name'])()
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        return getattr(tensors, settings[0]['name'])()
 
 
 class Abs(Transform):
+    """Returns :code:`abs(input)`"""
     def __init__(self, target: "Target" = 'update'): super().__init__({}, uses_grad=False, target=target)
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         torch._foreach_abs_(tensors)
         return tensors
 
 class Sign(Transform):
+    """Returns :code:`sign(input)`"""
     def __init__(self, target: "Target" = 'update'): super().__init__({}, uses_grad=False, target=target)
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         torch._foreach_sign_(tensors)
         return tensors
 
 class Exp(Transform):
+    """Returns :code:`exp(input)`"""
     def __init__(self, target: "Target" = 'update'): super().__init__({}, uses_grad=False, target=target)
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         torch._foreach_exp_(tensors)
         return tensors
 
 class Sqrt(Transform):
+    """Returns :code:`sqrt(input)`"""
     def __init__(self, target: "Target" = 'update'): super().__init__({}, uses_grad=False, target=target)
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         torch._foreach_sqrt_(tensors)
         return tensors
 
 class Reciprocal(Transform):
+    """Returns :code:`1 / input`"""
     def __init__(self, eps = 0, target: "Target" = 'update'):
         defaults = dict(eps = eps)
         super().__init__(defaults, uses_grad=False, target=target)
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        eps = self.get_settings('eps', params=params)
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        eps = [s['eps'] for s in settings]
         if any(e != 0 for e in eps): torch._foreach_add_(tensors, eps)
         torch._foreach_reciprocal_(tensors)
         return tensors
 
 class Negate(Transform):
+    """Returns :code:`- input`"""
     def __init__(self, target: "Target" = 'update'): super().__init__({}, uses_grad=False, target=target)
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         torch._foreach_neg_(tensors)
         return tensors
 
@@ -97,19 +113,19 @@ class NanToNum(Transform):
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        nan, posinf, neginf = self.get_settings('nan', 'posinf', 'neginf', params=params)
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        nan, posinf, neginf = unpack_dicts(settings, 'nan', 'posinf', 'neginf')
         return [t.nan_to_num_(nan_i, posinf_i, neginf_i) for t, nan_i, posinf_i, neginf_i in zip(tensors, nan, posinf, neginf)]
 
 class Rescale(Transform):
-    """rescale update to (min, max) range"""
+    """Rescales input to :code`(min, max)` range"""
     def __init__(self, min: float, max: float, tensorwise: bool = False, eps:float=1e-8, target: "Target" = 'update'):
         defaults = dict(min=min, max=max, eps=eps, tensorwise=tensorwise)
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        min,max = self.get_settings('min','max', params=params)
-        tensorwise = self.settings[params[0]]['tensorwise']
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        min, max = unpack_dicts(settings, 'min','max')
+        tensorwise = settings[0]['tensorwise']
         dim = None if tensorwise else 'global'
-        return TensorList(tensors).rescale(min=min, max=max, eps=self.settings[params[0]]['eps'], dim=dim)
+        return TensorList(tensors).rescale(min=min, max=max, eps=settings[0]['eps'], dim=dim)

torchzero/modules/ops/utility.py

@@ -6,107 +6,115 @@ from ...core import Module, Target, Transform
 from ...utils.tensorlist import Distributions, TensorList
 
 
-class Clone(Transform):
-    def __init__(self): super().__init__({}, uses_grad=False)
-    @torch.no_grad
-    def transform(self, tensors, params, grads, vars): return [t.clone() for t in tensors]
-
-class Grad(Module):
+class Clone(Module):
+    """Clones input. May be useful to store some intermediate result and make sure it doesn't get affected by in-place operations"""
     def __init__(self):
         super().__init__({})
     @torch.no_grad
-    def step(self, vars):
-        vars.update = [g.clone() for g in vars.get_grad()]
-        return vars
+    def step(self, var):
+        var.update = [u.clone() for u in var.get_update()]
+        return var
 
-class Params(Module):
+class Grad(Module):
+    """Outputs the gradient"""
     def __init__(self):
         super().__init__({})
     @torch.no_grad
-    def step(self, vars):
-        vars.update = [p.clone() for p in vars.params]
-        return vars
+    def step(self, var):
+        var.update = [g.clone() for g in var.get_grad()]
+        return var
 
-class Update(Module):
+class Params(Module):
+    """Outputs parameters"""
     def __init__(self):
         super().__init__({})
     @torch.no_grad
-    def step(self, vars):
-        vars.update = [u.clone() for u in vars.get_update()]
-        return vars
+    def step(self, var):
+        var.update = [p.clone() for p in var.params]
+        return var
 
 class Zeros(Module):
+    """Outputs zeros"""
     def __init__(self):
         super().__init__({})
     @torch.no_grad
-    def step(self, vars):
-        vars.update = [torch.zeros_like(p) for p in vars.params]
-        return vars
+    def step(self, var):
+        var.update = [torch.zeros_like(p) for p in var.params]
+        return var
 
 class Ones(Module):
+    """Outputs ones"""
     def __init__(self):
         super().__init__({})
     @torch.no_grad
-    def step(self, vars):
-        vars.update = [torch.ones_like(p) for p in vars.params]
-        return vars
+    def step(self, var):
+        var.update = [torch.ones_like(p) for p in var.params]
+        return var
 
 class Fill(Module):
+    """Outputs tensors filled with :code:`value`"""
     def __init__(self, value: float):
         defaults = dict(value=value)
         super().__init__(defaults)
 
     @torch.no_grad
-    def step(self, vars):
-        vars.update = [torch.full_like(p, self.settings[p]['value']) for p in vars.params]
-        return vars
+    def step(self, var):
+        var.update = [torch.full_like(p, self.settings[p]['value']) for p in var.params]
+        return var
 
 class RandomSample(Module):
+    """Outputs tensors filled with random numbers from distribution depending on value of :code:`distribution`."""
     def __init__(self, eps: float = 1, distribution: Distributions = 'normal'):
         defaults = dict(eps=eps, distribution=distribution)
         super().__init__(defaults)
 
     @torch.no_grad
-    def step(self, vars):
-        vars.update = TensorList(vars.params).sample_like(
-            eps=self.get_settings('eps',params=vars.params), distribution=self.settings[vars.params[0]]['distribution']
+    def step(self, var):
+        var.update = TensorList(var.params).sample_like(
+            eps=[self.settings[p]['eps'] for p in var.params], distribution=self.settings[var.params[0]]['distribution']
         )
-        return vars
+        return var
 
 class Randn(Module):
+    """Outputs tensors filled with random numbers from a normal distribution with mean 0 and variance 1."""
     def __init__(self):
         super().__init__({})
 
     @torch.no_grad
-    def step(self, vars):
-        vars.update = [torch.randn_like(p) for p in vars.params]
-        return vars
+    def step(self, var):
+        var.update = [torch.randn_like(p) for p in var.params]
+        return var
 
 class Uniform(Module):
+    """Outputs tensors filled with random numbers from uniform distribution between :code:`low` and :code:`high`."""
     def __init__(self, low: float, high: float):
         defaults = dict(low=low, high=high)
         super().__init__(defaults)
 
     @torch.no_grad
-    def step(self, vars):
-        low,high = self.get_settings('low','high', params=vars.params)
-        vars.update = [torch.empty_like(t).uniform_(l,h) for t,l,h in zip(vars.params, low, high)]
-        return vars
+    def step(self, var):
+        low,high = self.get_settings(var.params, 'low','high')
+        var.update = [torch.empty_like(t).uniform_(l,h) for t,l,h in zip(var.params, low, high)]
+        return var
 
 class GradToNone(Module):
+    """Sets :code:`grad` attribute to None on :code:`var`."""
     def __init__(self): super().__init__()
-    def step(self, vars):
-        vars.grad = None
-        return vars
+    def step(self, var):
+        var.grad = None
+        return var
 
 class UpdateToNone(Module):
+    """Sets :code:`update` attribute to None on :code:`var`."""
     def __init__(self): super().__init__()
-    def step(self, vars):
-        vars.update = None
-        return vars
+    def step(self, var):
+        var.update = None
+        return var
 
 class Identity(Module):
+    """A placeholder identity operator that is argument-insensitive."""
     def __init__(self, *args, **kwargs): super().__init__()
-    def step(self, vars): return vars
+    def step(self, var): return var
 
-NoOp = Identity
+NoOp = Identity
+"""A placeholder identity operator that is argument-insensitive."""

torchzero/modules/optimizers/__init__.py

@@ -1,7 +1,18 @@
 from .adagrad import Adagrad, FullMatrixAdagrad
+
+# from .curveball import CurveBall
+# from .spectral import SpectralPreconditioner
+from .adahessian import AdaHessian
 from .adam import Adam
+from .adan import Adan
+from .adaptive_heavyball import AdaptiveHeavyBall
+from .esgd import ESGD
+from .ladagrad import LMAdagrad
 from .lion import Lion
+from .mars import MARSCorrection
+from .msam import MSAM, MSAMObjective
 from .muon import DualNormCorrection, MuonAdjustLR, Orthogonalize, orthogonalize_grads_
+from .orthograd import OrthoGrad, orthograd_
 from .rmsprop import RMSprop
 from .rprop import (
     BacktrackOnSignChange,
@@ -10,9 +21,7 @@ from .rprop import (
     SignConsistencyLRs,
     SignConsistencyMask,
 )
+from .sam import ASAM, SAM
 from .shampoo import Shampoo
 from .soap import SOAP
-from .orthograd import OrthoGrad, orthograd_
 from .sophia_h import SophiaH
-# from .curveball import CurveBall
-# from .spectral import SpectralPreconditioner

torchzero/modules/optimizers/adagrad.py

@@ -1,18 +1,17 @@
 from operator import itemgetter
+from typing import Literal
 
 import torch
-
 from ...core import (
     Chainable,
     Module,
-    Preconditioner,
     Target,
-    TensorwisePreconditioner,
+    TensorwiseTransform,
     Transform,
-    Vars,
-    apply,
+    Var,
+    apply_transform,
 )
-from ...utils import NumberList, TensorList
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
 from ...utils.linalg import matrix_power_eigh
 from ..functional import add_power_, lerp_power_, root
 
@@ -26,12 +25,12 @@ def adagrad_(
     step: int,
     pow: float = 2,
     use_sqrt: bool = True,
+    divide: bool = False,
 
     # inner args
     inner: Module | None = None,
     params: list[torch.Tensor] | None = None,
     grads: list[torch.Tensor] | None = None,
-    vars: Vars | None = None,
 ):
     """returns `tensors_`"""
     clr = alpha / (1 + step * lr_decay)
@@ -40,7 +39,9 @@ def adagrad_(
 
     if inner is not None:
         assert params is not None
-        tensors_ = TensorList(apply(inner, tensors_, params=params, grads=grads, vars=vars))
+        tensors_ = TensorList(apply_transform(inner, tensors_, params=params, grads=grads))
+
+    if divide: sq_sum_ = sq_sum_ / max(step, 1)
 
     if use_sqrt: tensors_.div_(root(sq_sum_, p=pow, inplace=False).add_(eps)).mul_(clr)
     else: tensors_.div_(sq_sum_.add(eps)).mul_(clr)
@@ -50,7 +51,9 @@ def adagrad_(
 
 
 class Adagrad(Transform):
-    """Adagrad, divides by sum of past squares of gradients, matches pytorch Adagrad.
+    """Adagrad, divides by sum of past squares of gradients.
+
+    This implementation is identical to :code:`torch.optim.Adagrad`.
 
     Args:
         lr_decay (float, optional): learning rate decay. Defaults to 0.
@@ -69,29 +72,30 @@ class Adagrad(Transform):
         alpha: float = 1,
         pow: float = 2,
         use_sqrt: bool = True,
+        divide: bool=False,
         inner: Chainable | None = None,
     ):
         defaults = dict(alpha = alpha, lr_decay = lr_decay, initial_accumulator_value=initial_accumulator_value,
-                        eps = eps, pow=pow, use_sqrt = use_sqrt)
+                        eps = eps, pow=pow, use_sqrt = use_sqrt, divide=divide)
         super().__init__(defaults=defaults, uses_grad=False)
 
         if inner is not None:
             self.set_child('inner', inner)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         tensors = TensorList(tensors)
         step = self.global_state['step'] = self.global_state.get('step', 0) + 1
 
-        lr_decay,alpha,eps = self.get_settings('lr_decay', 'alpha', 'eps', params=params, cls=NumberList)
+        lr_decay,alpha,eps = unpack_dicts(settings, 'lr_decay', 'alpha', 'eps', cls=NumberList)
 
-        pow, use_sqrt = itemgetter('pow', 'use_sqrt')(self.settings[params[0]])
+        pow, use_sqrt, divide = itemgetter('pow', 'use_sqrt', 'divide')(settings[0])
 
-        sq_sum = self.get_state('sq_sum', params=params, cls=TensorList)
+        sq_sum = unpack_states(states, tensors, 'sq_sum', cls=TensorList)
 
         # initialize accumulator on 1st step
         if step == 1:
-            sq_sum.set_(tensors.full_like(self.get_settings('initial_accumulator_value', params=params)))
+            sq_sum.set_(tensors.full_like([s['initial_accumulator_value'] for s in settings]))
 
         return adagrad_(
             tensors,
@@ -102,45 +106,60 @@ class Adagrad(Transform):
             step=self.global_state["step"],
             pow=pow,
             use_sqrt=use_sqrt,
+            divide=divide,
 
             # inner args
             inner=self.children.get("inner", None),
             params=params,
             grads=grads,
-            vars=vars,
         )
 
 
 
-class FullMatrixAdagrad(TensorwisePreconditioner):
-    def __init__(self, beta: float | None = None, decay: float | None = None, concat_params=False, update_freq=1, inner: Chainable | None = None):
-        defaults = dict(beta=beta, decay=decay)
-        super().__init__(defaults, uses_grad=False, concat_params=concat_params, update_freq=update_freq, inner=inner)
+class FullMatrixAdagrad(TensorwiseTransform):
+    def __init__(self, beta: float | None = None, decay: float | None = None, sqrt:bool=True, concat_params=True, update_freq=1, init: Literal['identity', 'zeros', 'ones', 'GGT'] = 'identity', divide: bool=False, inner: Chainable | None = None):
+        defaults = dict(beta=beta, decay=decay, sqrt=sqrt, init=init, divide=divide)
+        super().__init__(defaults, uses_grad=False, concat_params=concat_params, update_freq=update_freq, inner=inner,)
 
     @torch.no_grad
-    def update_tensor(self, tensor, param, grad, state, settings):
+    def update_tensor(self, tensor, param, grad, loss, state, setting):
         G = tensor.ravel()
         GG = torch.outer(G, G)
-        decay = settings['decay']
-        beta = settings['beta']
-
-        if 'GG' not in state: state['GG'] = torch.eye(GG.size(0), device=GG.device, dtype=GG.dtype)
+        decay = setting['decay']
+        beta = setting['beta']
+        init = setting['init']
+
+        if 'GG' not in state:
+            if init == 'identity': state['GG'] = torch.eye(GG.size(0), device=GG.device, dtype=GG.dtype)
+            elif init == 'zeros': state['GG'] =  torch.zeros_like(GG)
+            elif init == 'ones': state['GG'] = torch.ones_like(GG)
+            elif init == 'GGT': state['GG'] = GG.clone()
+            else: raise ValueError(init)
         if decay is not None: state['GG'].mul_(decay)
 
         if beta is not None: state['GG'].lerp_(GG, 1-beta)
         else: state['GG'].add_(GG)
+        state['i'] = state.get('i', 0) + 1 # number of GGTs in sum
 
     @torch.no_grad
-    def apply_tensor(self, tensor, param, grad, state, settings):
+    def apply_tensor(self, tensor, param, grad, loss, state, setting):
         GG = state['GG']
+        sqrt = setting['sqrt']
+        divide = setting['divide']
+        if divide: GG = GG/state.get('i', 1)
 
         if tensor.numel() == 1:
-            return tensor / (GG**(1/2)).squeeze()
+            GG = GG.squeeze()
+            if sqrt: return tensor / GG.sqrt()
+            return tensor / GG
 
         try:
-            B = matrix_power_eigh(GG, -1/2)
+            if sqrt: B = matrix_power_eigh(GG, -1/2)
+            else: return torch.linalg.solve(GG, tensor.ravel()).view_as(tensor) # pylint:disable = not-callable
+
         except torch.linalg.LinAlgError:
-            return tensor.div_(tensor.abs().max()) # conservative scaling
+            scale = 1 / tensor.abs().max()
+            return tensor.mul_(scale.clip(min=torch.finfo(tensor.dtype).eps, max=1)) # conservative scaling
 
         return (B @ tensor.ravel()).view_as(tensor)