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

torchzero/utils/optimizer.py

@@ -1,3 +1,4 @@
+from abc import ABC, abstractmethod
 from collections.abc import Callable, Iterable, Mapping, MutableSequence, Sequence, MutableMapping
 from typing import Any, Literal, TypeVar, overload
 
@@ -132,65 +133,7 @@ def get_state_vals(state: Mapping[torch.Tensor, MutableMapping[str, Any]], param
     return values
 
 
-
-def loss_at_params(closure, params: Iterable[torch.Tensor],
-                   new_params: Sequence[torch.Tensor] | Any, backward: bool, restore=False):
-    params = TensorList(params)
-
-    old_params = params.clone() if restore else None
-
-    if isinstance(new_params, Sequence) and isinstance(new_params[0], torch.Tensor):
-        # when not restoring, copy new_params to params to avoid unexpected bugs due to shared storage
-        # when restoring params will be set back to old_params so its fine
-        if restore: params.set_(new_params)
-        else: params.copy_(new_params) # type:ignore
-
-    else:
-        new_params = totensor(new_params)
-        params.from_vec_(new_params)
-
-    if backward: loss = closure()
-    else: loss = closure(False)
-
-    if restore:
-        assert old_params is not None
-        params.set_(old_params)
-
-    return tofloat(loss)
-
-def loss_grad_at_params(closure, params: Iterable[torch.Tensor], new_params: Sequence[torch.Tensor], restore=False):
-    params = TensorList(params)
-    old_params = params.clone() if restore else None
-    loss = loss_at_params(closure, params, new_params, backward=True, restore=False)
-    grad = params.ensure_grad_().grad
-
-    if restore:
-        assert old_params is not None
-        params.set_(old_params)
-
-    return loss, grad
-
-def grad_at_params(closure, params: Iterable[torch.Tensor], new_params: Sequence[torch.Tensor], restore=False):
-    return loss_grad_at_params(closure=closure,params=params,new_params=new_params,restore=restore)[1]
-
-def loss_grad_vec_at_params(closure, params: Iterable[torch.Tensor], new_params: Any, restore=False):
-    params = TensorList(params)
-    old_params = params.clone() if restore else None
-    loss = loss_at_params(closure, params, new_params, backward=True, restore=False)
-    grad = params.ensure_grad_().grad.to_vec()
-
-    if restore:
-        assert old_params is not None
-        params.set_(old_params)
-
-    return loss, grad
-
-def grad_vec_at_params(closure, params: Iterable[torch.Tensor], new_params: Any, restore=False):
-    return loss_grad_vec_at_params(closure=closure,params=params,new_params=new_params,restore=restore)[1]
-
-
-
-class Optimizer(torch.optim.Optimizer):
+class Optimizer(torch.optim.Optimizer, ABC):
     """subclass of torch.optim.Optimizer with some helper methods for fast experimentation, it's not used anywhere in torchzero.
 
     Args:
@@ -251,21 +194,10 @@ class Optimizer(torch.optim.Optimizer):
 
         return get_state_vals(self.state, params, key, key2, *keys, init = init, cls = cls) # type:ignore[reportArgumentType]
 
-    def loss_at_params(self, closure, params: Sequence[torch.Tensor] | Any, backward: bool, restore=False):
-        return loss_at_params(closure=closure,params=self.get_params(),new_params=params,backward=backward,restore=restore)
-
-    def loss_grad_at_params(self, closure, params: Sequence[torch.Tensor] | Any, restore=False):
-        return loss_grad_at_params(closure=closure,params=self.get_params(),new_params=params,restore=restore)
-
-    def grad_at_params(self, closure, new_params: Sequence[torch.Tensor], restore=False):
-        return self.loss_grad_at_params(closure=closure,params=new_params,restore=restore)[1]
-
-    def loss_grad_vec_at_params(self, closure, params: Any, restore=False):
-        return loss_grad_vec_at_params(closure=closure,params=self.get_params(),new_params=params,restore=restore)
-
-    def grad_vec_at_params(self, closure, params: Any, restore=False):
-        return self.loss_grad_vec_at_params(closure=closure,params=params,restore=restore)[1]
 
+    # shut up pylance
+    @abstractmethod
+    def step(self, closure) -> Any: ... # pylint:disable=signature-differs # pyright:ignore[reportIncompatibleMethodOverride]
 
 def zero_grad_(params: Iterable[torch.Tensor], set_to_none):
     if set_to_none:
@@ -281,4 +213,53 @@ def zero_grad_(params: Iterable[torch.Tensor], set_to_none):
             else:
                 grad.requires_grad_(False)
 
-        torch._foreach_zero_(grads)
+        torch._foreach_zero_(grads)
+
+
+@overload
+def unpack_states(states: Sequence[MutableMapping[str, Any]], tensors: Sequence[torch.Tensor],
+                   key: str, *,
+                   must_exist: bool = False, init: Init = torch.zeros_like,
+                   cls: type[ListLike] = list) -> ListLike: ...
+@overload
+def unpack_states(states: Sequence[MutableMapping[str, Any]], tensors: Sequence[torch.Tensor],
+                   key: list[str] | tuple[str,...], *,
+                   must_exist: bool = False, init: Init | Sequence[Init] = torch.zeros_like,
+                   cls: type[ListLike] = list) -> list[ListLike]: ...
+@overload
+def unpack_states(states: Sequence[MutableMapping[str, Any]], tensors: Sequence[torch.Tensor],
+                   key: str,  key2: str, *keys: str,
+                   must_exist: bool = False, init: Init | Sequence[Init] = torch.zeros_like,
+                   cls: type[ListLike] = list) -> list[ListLike]: ...
+
+def unpack_states(states: Sequence[MutableMapping[str, Any]], tensors: Sequence[torch.Tensor],
+                   key: str | list[str] | tuple[str,...], key2: str | None = None,  *keys: str,
+                   must_exist: bool = False, init: Init | Sequence[Init] = torch.zeros_like,
+                   cls: type[ListLike] = list) -> ListLike | list[ListLike]:
+
+    # single key, return single cls
+    if isinstance(key, str) and key2 is None:
+        values = cls()
+        for i,s in enumerate(states):
+            if key not in s:
+                if must_exist: raise KeyError(f"Key {key} doesn't exist in state with keys {tuple(s.keys())}")
+                s[key] = _make_initial_state_value(tensors[i], init, i)
+            values.append(s[key])
+        return values
+
+    # multiple keys
+    k1 = (key,) if isinstance(key, str) else tuple(key)
+    k2 = () if key2 is None else (key2,)
+    keys = k1 + k2 + keys
+
+    values = [cls() for _ in keys]
+    for i,s in enumerate(states):
+        for k_i, key in enumerate(keys):
+            if key not in s:
+                if must_exist: raise KeyError(f"Key {key} doesn't exist in state with keys {tuple(s.keys())}")
+                k_init = init[k_i] if isinstance(init, (list,tuple)) else init
+                s[key] = _make_initial_state_value(tensors[i], k_init, i)
+            values[k_i].append(s[key])
+
+    return values
+

torchzero/utils/python_tools.py

@@ -1,7 +1,7 @@
 import functools
 import operator
-from typing import Any, TypeVar
-from collections.abc import Iterable, Callable
+from typing import Any, TypeVar, overload
+from collections.abc import Iterable, Callable, Mapping, MutableSequence
 from collections import UserDict
 
 
@@ -17,8 +17,8 @@ def flatten(iterable: Iterable) -> list[Any]:
     raise TypeError(f'passed object is not an iterable, {type(iterable) = }')
 
 X = TypeVar("X")
-# def reduce_dim[X](x:Iterable[Iterable[X]]) -> list[X]: # pylint:disable=E0602
-def reduce_dim(x:Iterable[Iterable[X]]) -> list[X]: # pylint:disable=E0602
+# def reduce_dim[X](x:Iterable[Iterable[X]]) -> list[X]:
+def reduce_dim(x:Iterable[Iterable[X]]) -> list[X]:
     """Reduces one level of nesting. Takes an iterable of iterables of X, and returns an iterable of X."""
     return functools.reduce(operator.iconcat, x, [])
 
@@ -38,3 +38,16 @@ def zipmap(self, fn: Callable, other: Any | list | tuple, *args, **kwargs):
     if isinstance(other, (list, tuple)): return self.__class__(fn(i, j, *args, **kwargs) for i, j in zip(self, other))
     return self.__class__(fn(i, other, *args, **kwargs) for i in self)
 
+ListLike = TypeVar('ListLike', bound=MutableSequence)
+@overload
+def unpack_dicts(dicts: Iterable[Mapping[str, Any]], key:str, *, cls:type[ListLike]=list) -> ListLike: ...
+@overload
+def unpack_dicts(dicts: Iterable[Mapping[str, Any]], key:str, key2: str, *keys:str, cls:type[ListLike]=list) -> list[ListLike]: ...
+def unpack_dicts(dicts: Iterable[Mapping[str, Any]], key:str, key2: str | None = None, *keys:str, cls:type[ListLike]=list) -> ListLike | list[ListLike]:
+    k1 = (key,) if isinstance(key, str) else tuple(key)
+    k2 = () if key2 is None else (key2,)
+    keys = k1 + k2 + keys
+
+    values = [cls(s[k] for s in dicts) for k in keys] # pyright:ignore[reportCallIssue]
+    if len(values) == 1: return values[0]
+    return values

{torchzero-0.3.9.dist-info → torchzero-0.3.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchzero
-Version: 0.3.9
+Version: 0.3.10
 Summary: Modular optimization library for PyTorch.
 Author-email: Ivan Nikishev <nkshv2@gmail.com>
 License: MIT License
@@ -157,13 +157,14 @@ for epoch in range(100):
   * `NewtonCG`: Matrix-free newton's method with conjugate gradient solver.
   * `NystromSketchAndSolve`: Nyström sketch-and-solve method.
   * `NystromPCG`: NewtonCG with Nyström preconditioning (usually beats NewtonCG).
+  * `HigherOrderNewton`: Higher order Newton's method with trust region.
 
 * **Quasi-Newton**: Approximate second-order optimization methods.
   * `LBFGS`: Limited-memory BFGS.
   * `LSR1`: Limited-memory SR1.
   * `OnlineLBFGS`: Online LBFGS.
-  * `BFGS`, `SR1`, `DFP`, `BroydenGood`, `BroydenBad`, `Greenstadt1`, `Greenstadt2`, `ColumnUpdatingMethod`, `ThomasOptimalMethod`, `PSB`, `Pearson2`, `SSVM`: Classic full-matrix quasi-newton methods.
-  * `PolakRibiere`, `FletcherReeves`, `HestenesStiefel`, `DaiYuan`, `LiuStorey`, `ConjugateDescent`, `HagerZhang`, `HybridHS_DY`: Conjugate gradient methods.
+  * `BFGS`, `DFP`, `PSB`, `SR1`, `SSVM`, `BroydenBad`, `BroydenGood`, `ColumnUpdatingMethod`, `FletcherVMM`, `GradientCorrection`, `Greenstadt1`, `Greenstadt2`, `Horisho`, `McCormick`, `Pearson`, `ProjectedNewtonRaphson`, `ThomasOptimalMethod`: Classic full-matrix quasi-newton methods.
+  * `PolakRibiere`, `FletcherReeves`, `HestenesStiefel`, `DaiYuan`, `LiuStorey`, `ConjugateDescent`, `HagerZhang`, `HybridHS_DY`, `ProjectedGradientMethod`: Conjugate gradient methods.
 
 * **Line Search**:
   * `Backtracking`, `AdaptiveBacktracking`: Backtracking line searches (adaptive is my own).
@@ -312,20 +313,20 @@ not in the module itself. Also both per-parameter settings and state are stored
 
 ```python
 import torch
-from torchzero.core import Module, Vars
+from torchzero.core import Module, Var
 
 class HeavyBall(Module):
     def __init__(self, momentum: float = 0.9, dampening: float = 0):
         defaults = dict(momentum=momentum, dampening=dampening)
         super().__init__(defaults)
 
-    def step(self, vars: Vars):
-        # a module takes a Vars object, modifies it or creates a new one, and returns it
-        # Vars has a bunch of attributes, including parameters, gradients, update, closure, loss
+    def step(self, var: Var):
+        # a module takes a Var object, modifies it or creates a new one, and returns it
+        # Var has a bunch of attributes, including parameters, gradients, update, closure, loss
         # for now we are only interested in update, and we will apply the heavyball rule to it.
 
-        params = vars.params
-        update = vars.get_update() # list of tensors
+        params = var.params
+        update = var.get_update() # list of tensors
 
         exp_avg_list = []
         for p, u in zip(params, update):
@@ -346,16 +347,15 @@ class HeavyBall(Module):
             # and it is part of self.state
             exp_avg_list.append(buf.clone())
 
-        # set new update to vars
-        vars.update = exp_avg_list
-        return vars
+        # set new update to var
+        var.update = exp_avg_list
+        return var
 ```
 
 There are a some specialized base modules that make it much easier to implement some specific things.
 
 * `GradApproximator` for gradient approximations
 * `LineSearch` for line searches
-* `Preconditioner` for preconditioners
 * `Projection` for projections like GaLore or into fourier domain.
 * `QuasiNewtonH` for full-matrix quasi-newton methods that update hessian inverse approximation (because they are all very similar)
 * `ConguateGradientBase` for conjugate gradient methods, basically the only difference is how beta is calculated.
@@ -376,4 +376,4 @@ There are also wrappers providing `torch.optim.Optimizer` interface for for `sci
 
 They are in `torchzero.optim.wrappers.scipy.ScipyMinimize`, `torchzero.optim.wrappers.nlopt.NLOptOptimizer`, and `torchzero.optim.wrappers.nevergrad.NevergradOptimizer`. Make sure closure has `backward` argument as described in **Advanced Usage**.
 
-Apparently https://github.com/avaneev/biteopt is diabolical so I will add a wrapper for it too very soon.
+Apparently <https://github.com/avaneev/biteopt> is diabolical so I will add a wrapper for it too very soon.

torchzero-0.3.10.dist-info/RECORD

@@ -0,0 +1,139 @@
+docs/source/conf.py,sha256=jd80ZT2IdCx7nlQrpOTJL8UhGBNm6KYyXlpp0jmRiAw,1849
+tests/test_identical.py,sha256=NZ7A8Rm1U9Q16d-cG2G_wccpPtNALyoKYJt9qMownMc,11568
+tests/test_module.py,sha256=qX3rjdSJsbA8JO17bPTUIDspe7bg2dogqxMw__KV7SU,2039
+tests/test_opts.py,sha256=VSko5fUuACo_y6iab_akke0gMhCUEEUJ9ahpBqWBoM4,41715
+tests/test_tensorlist.py,sha256=SwzLKLrs2ppMtm_7UrfTDTlD-ObZd7JQ_FNHbp059tc,72460
+tests/test_utils_optimizer.py,sha256=bvC0Ehvs2L8fohpyIF5Vfr9OKTycpnODWLPflXilU1c,8414
+tests/test_vars.py,sha256=MqCJXrbj-C75APm1heykzcEWewinihlSjekkYDx-TFk,6726
+torchzero/__init__.py,sha256=L7IJ1qZ3o8E9oRwlJZBK2_2yII_eeGEk57Of6EfVbrk,112
+torchzero/core/__init__.py,sha256=Zib_4is13LFAabp_7VU8QXZpQEEZGzsH94vgRI0HxAg,150
+torchzero/core/module.py,sha256=Yfzn48dDbxYZJLpWnLYFIbqBb4sB3GekSZ7QGIplYAg,27525
+torchzero/core/transform.py,sha256=yK1wYgp03THzRN9y_f9-5q2nonEZMa0CfDFAdOxnqEU,11778
+torchzero/modules/__init__.py,sha256=8C73_dFzfWUWhii1UF86FUy8x75RPiAVLAm4sLTikBg,359
+torchzero/modules/functional.py,sha256=HXNzmPe7LsPadryEm7zrcEKqGej16QDwSgBkbEvggFM,6492
+torchzero/modules/clipping/__init__.py,sha256=ZaffMF7mIRK6hZSfuZadgjNTX6hF5ANiLBny2w3S7I8,250
+torchzero/modules/clipping/clipping.py,sha256=XKFKvzNgsvuYUmvHyulE6PkZv_aeLQjp0CgtFj0013s,12516
+torchzero/modules/clipping/ema_clipping.py,sha256=MGouZEN0BorliHAZhue0afhC3AhZJ6wrnwBRzDTHjX4,5978
+torchzero/modules/clipping/growth_clipping.py,sha256=50c1YOUPVL8eWzH6zJINnNP68oiZkDcq7rR6HnWjVFc,6674
+torchzero/modules/experimental/__init__.py,sha256=zxxNKPZHnkVnx1ZjKNX_nkV4Wc_EdODM6qJGn7Pgb3w,766
+torchzero/modules/experimental/absoap.py,sha256=-KwQXmI12hvHbMGPHM0APAxDQztlFhlSOG55KK6PvpI,9901
+torchzero/modules/experimental/adadam.py,sha256=o0KPLaF4J7L_Ty71RNgsysk6IEuC4DRE5nGQkGIP_dA,4078
+torchzero/modules/experimental/adamY.py,sha256=LZabWX_vccDaG6_UVZl9ALJ-3nCZu-NEygJQ_Bwzel8,4018
+torchzero/modules/experimental/adasoap.py,sha256=XtxEvBWYdcqfWnQqOFa_-SrOwd_nXHzLftiw-YXDACQ,7408
+torchzero/modules/experimental/curveball.py,sha256=JdgojuSYLNe9u3bmqcYrFm8brUD4kvKm9XYx78GzpKI,3257
+torchzero/modules/experimental/diagonal_higher_order_newton.py,sha256=u4-a5qJ_97XiZUDlClE2cASkBsx_NTJNPk6BWWybiqE,7158
+torchzero/modules/experimental/eigendescent.py,sha256=0cM1p4rYbrpwBNXgBEMblVyX0xBWTzojSC1EsUnXH6k,4707
+torchzero/modules/experimental/etf.py,sha256=FsLOCmQf24PPoRf5wsRUjVqk32uW9uTzaf1ERjFxRK8,5744
+torchzero/modules/experimental/gradmin.py,sha256=UixSLdca4ekYHOipEivdXfBAV-uEL9TZm5nCFXVaNco,3684
+torchzero/modules/experimental/newton_solver.py,sha256=3dZ7FG-2vGxJKkFF9P2LCs-LI_epcvZbyNtJOtw47pg,3055
+torchzero/modules/experimental/newtonnewton.py,sha256=QCGnY_CFo0i_NUB7D-6ezeNpG6wLkTD5lHBiakFIqbM,3033
+torchzero/modules/experimental/reduce_outward_lr.py,sha256=VFjcTpmLwpfhUR8u_5rbzPgHVR6K3fvti7jVy1DnsYU,1300
+torchzero/modules/experimental/soapy.py,sha256=7qsh9Y9U9oeQDwuDSVqnz71AD0nUYY3q0XN2XoMFWaw,6721
+torchzero/modules/experimental/spectral.py,sha256=SN7tToIpmna0IZ1NgObvqEbO48NnVbwqRwKi8ROsb7s,7374
+torchzero/modules/experimental/structured_newton.py,sha256=CWfVJ2LPZUuz1bMnlgOM6tlYPd2etjgLDIcyAfAG_y8,3464
+torchzero/modules/experimental/subspace_preconditioners.py,sha256=9Tl1PCN9crFUvVn6343GHoI3kv6CVnUWP1dfhwUvAFU,5130
+torchzero/modules/experimental/tada.py,sha256=84YcLhG34CbWq84L-AUj-A4uxpzdIVayaARHRm2f9b8,1564
+torchzero/modules/grad_approximation/__init__.py,sha256=DVFjf0cXuF70NA0nJ2WklpP01PQgrRZxUjUQjjQeSos,195
+torchzero/modules/grad_approximation/fdm.py,sha256=cUgy98Bz0Br4q6ViNxn6EVOZX2jE0nDXVZLUGhxpDcA,3589
+torchzero/modules/grad_approximation/forward_gradient.py,sha256=cNgx8kc8r0fWj8xdU2b85W3fenNDQZKuIsJLM3UzSig,3867
+torchzero/modules/grad_approximation/grad_approximator.py,sha256=TODFUwBgTmjfbnO6Sc833fnvYzYaqqYTEba_13s-qOI,2906
+torchzero/modules/grad_approximation/rfdm.py,sha256=VsRlf95JnG6HdlIsJANcfJjMk7c_B9a5-fH9dSTBA10,11328
+torchzero/modules/higher_order/__init__.py,sha256=W94CY8K1NFxs9TPi415UssKVKz5MV_bH9adax1uZsYM,50
+torchzero/modules/higher_order/higher_order_newton.py,sha256=BwiSlcGobam04SgWFcB1p_-TSuzu2rWgGVnmvP6Si9k,9751
+torchzero/modules/line_search/__init__.py,sha256=nkOUPLe88wE91ICEhprl2pJsvaKtbI3KzYOdT83AGsg,253
+torchzero/modules/line_search/backtracking.py,sha256=ZgeLAYqrw-6BeEGp8wmOgFoLtUKROF7w7LpAREe0xZU,7704
+torchzero/modules/line_search/line_search.py,sha256=CfOENZgAPSdyv1wvSbhw6gdpfbQnXGdOnLsq29wjvzU,7229
+torchzero/modules/line_search/scipy.py,sha256=SvDCZ1DPOLZcSeOFvf3tXAf1ty-9qRVfGFMWVF5q708,2293
+torchzero/modules/line_search/strong_wolfe.py,sha256=xOU4XFekh4TIepm9ztJTYpcGucEMPwAeb_cDK4Rp0ho,7620
+torchzero/modules/line_search/trust_region.py,sha256=xUZApOTW4uXFBk_Uq_YBktiXcoSAKdDc6O5vjTwquGw,3101
+torchzero/modules/lr/__init__.py,sha256=kh2k_tma-oTOALR6AlD5XHdTPSMgU4A04Oa0hAqrEpI,89
+torchzero/modules/lr/adaptive.py,sha256=6s06Gvu1UmoT89hrMkXvJWHkEOMNcy5mMiyxy3V9lQs,3904
+torchzero/modules/lr/lr.py,sha256=1gU2QzMA5PV2KkzOkxxrZZKGcz-Kbjyp7WNurOM36ys,2655
+torchzero/modules/momentum/__init__.py,sha256=pSD7vxu8PySrYOSHQMi3C9heYdcQr8y6WC_rwMybZm0,544
+torchzero/modules/momentum/averaging.py,sha256=NmRodxsSekEDGIuFGDYOvJL-WkdMN3YF-naBdtfjxx8,3247
+torchzero/modules/momentum/cautious.py,sha256=JuaFYfyf9S3nTcqeZz5ylXKepqi0eqglOAQ0uNG0eT8,7373
+torchzero/modules/momentum/ema.py,sha256=qJV__nIbcD9e8qvwbvsATnYkQrdnmMiA91ju52IqSxw,10699
+torchzero/modules/momentum/experimental.py,sha256=eYnP6NmBDegwX9XC_dYMJP3vquBpM1LyQc03v3vW6-8,6900
+torchzero/modules/momentum/matrix_momentum.py,sha256=LR12UugXM8ocwTB8zBYpt03oZeZU0cb0UoFR6qO34V8,6818
+torchzero/modules/momentum/momentum.py,sha256=4Pgk-3HM7Av_ILT6oXtvnM1CB1yit8AkFnYWLvnUAqs,2655
+torchzero/modules/ops/__init__.py,sha256=hxMZFSXX7xvitXkuBiYykVGX3p03Xprm_QA2CMg4eW8,1601
+torchzero/modules/ops/accumulate.py,sha256=yKNgw8ZsaVRPjuzPzLJOvALkjik0aWx30Eu91FefRoA,3741
+torchzero/modules/ops/binary.py,sha256=98jyjkJ8BPuSH-mb4g2BnFi6UzvRZRf__Pt-jnD3pNU,9690
+torchzero/modules/ops/debug.py,sha256=zueWyNVvaJmxRD8QG8m_ys9jc7zRfSr8GAuxqz5dDTI,851
+torchzero/modules/ops/misc.py,sha256=GmnKDjMXaTUjPcC5e7Jftk6k2NQ0Ivv4ceUApxMckIQ,15978
+torchzero/modules/ops/multi.py,sha256=T1aVaRr6bLWvjoj1cyxaDncROypT6rmmmji8mvBHczo,5713
+torchzero/modules/ops/reduce.py,sha256=reGvusJyCzM8VdHbWyJRYFePPBXfVP0jZeXIEKGIJGc,5668
+torchzero/modules/ops/split.py,sha256=eM4Qsz6pMNF22bk3NF2rtvyxSOt9U-EyYxMAyjvTrMQ,2265
+torchzero/modules/ops/switch.py,sha256=ddsxq4bsH86iWW6mMdcQw3c0mU1s2FA-PRZpVOia7PY,2506
+torchzero/modules/ops/unary.py,sha256=3ysDHXs6snsQNBj3c288BT8G6T30Nvo0QM3PcdfQ2ww,4888
+torchzero/modules/ops/utility.py,sha256=8XFjQO4ghCmGD2H-lYTgaBzik_9pB3Uxt7xCxQrv5Ig,3080
+torchzero/modules/optimizers/__init__.py,sha256=BbT2nhIt4p74t1cO8ziQgzqZHaLvyuleXQbccugd06M,554
+torchzero/modules/optimizers/adagrad.py,sha256=NHpWcnIRM2LyPnNtDVTdluX4n1qqqWs9IHpFD8uYeLo,5500
+torchzero/modules/optimizers/adam.py,sha256=u6ieXHn_lHZozmGiKhSA73pApI83eeTNIyOrxBTFL1o,4009
+torchzero/modules/optimizers/lion.py,sha256=4yy6d0SLpGXndu8NCuYhdsNshMEYhONu_FPYXdupA_s,1119
+torchzero/modules/optimizers/muon.py,sha256=exbp7wVpIryiOxmbf9RAfZ9a6XXuOWTUqdjn-i57Fq4,9628
+torchzero/modules/optimizers/orthograd.py,sha256=cN5g7OusfeUlh0jn0kjkvpcVjqR01eGoi9WK1sSPnug,2021
+torchzero/modules/optimizers/rmsprop.py,sha256=jM5ohfABYUljy2RrtG_bY9PMHNzSkROYjqFPxnlXE6o,4309
+torchzero/modules/optimizers/rprop.py,sha256=d0R8UR-f9Pb64VrsJegrCyteLYa5TAmgObjgirqLaBo,11030
+torchzero/modules/optimizers/shampoo.py,sha256=hmfgPghzmjmba3PH1vLzaz0lOvLiIX9rCKrT71YZb40,8420
+torchzero/modules/optimizers/soap.py,sha256=7adybqncrkt31rNveQwXp8eeZKWf0LDhC5wt7GbmDcM,11052
+torchzero/modules/optimizers/sophia_h.py,sha256=He9YrHeaQhiz4CJm-3H_d_M07MGTsP663v8wx4BlaZI,4273
+torchzero/modules/projections/__init__.py,sha256=OCxlh_-Tx-xpl31X03CeFJH9XveH563oEsWc8rUvX0A,196
+torchzero/modules/projections/dct.py,sha256=0tswjgta3mE5D5Yjw9mJWqPBDga0OIe3lKlwd1AXASc,2369
+torchzero/modules/projections/fft.py,sha256=wNDZP5-3b2-bND3qH2yvX3SqFaljbLkPTQ1gUnlH5fU,2955
+torchzero/modules/projections/galore.py,sha256=etaG2gxazxuDEu-e2r7lKIKMTPEGGS5Vi7HXccmD3kY,241
+torchzero/modules/projections/projection.py,sha256=QUV_Gi6QlPiWEmcc7rwucr2yuYwYFGvSRUAT4uucqMY,10049
+torchzero/modules/projections/structural.py,sha256=f8-72zViXJ6S2gxDagkrrul9IaOPsYXZmX8sFLYkxCc,5683
+torchzero/modules/quasi_newton/__init__.py,sha256=Yc-NV__cJCiYLr2BZG4VsYa3VVq4gCxBMcirQEXSNIo,630
+torchzero/modules/quasi_newton/cg.py,sha256=lvmwJNTR7AEcpDIvpcLnMrZrOLwNld8GFAC19CcTKoY,11661
+torchzero/modules/quasi_newton/lbfgs.py,sha256=BDiv3f7qN8-Nhs8LqtWwk7Wwv88NtXXYle5WwKeekm4,9198
+torchzero/modules/quasi_newton/lsr1.py,sha256=A0Pstikb6JrQbwM5RZjLw9WJEHiMRy3PsPF1_iLkrK4,6053
+torchzero/modules/quasi_newton/olbfgs.py,sha256=Tz2eubiN7OXGN1mbXT4VKPd9kynpXzcLas7mrvBax-k,8333
+torchzero/modules/quasi_newton/quasi_newton.py,sha256=4hRII9GFE5MzNtXkHH_T1hEJ1T8T4-Q4A4MXlhf64mc,25142
+torchzero/modules/quasi_newton/experimental/__init__.py,sha256=3qpZGgdsx6wpoafWaNWx-eamRl1FuxVCWQZq8Y7Cl98,39
+torchzero/modules/quasi_newton/experimental/modular_lbfgs.py,sha256=oLbJ96sl-2XBwLbJrrTZiLJIhKhTPOD6-wny7hbSno4,10767
+torchzero/modules/second_order/__init__.py,sha256=jolCGaIVkID9hpxgx0Tc22wgjVlwuWekWjKTMe5jKXw,114
+torchzero/modules/second_order/newton.py,sha256=ZYIcLpifcOHL_KRC6YoNs-MJQKM39urXUQzReWnWeXE,6583
+torchzero/modules/second_order/newton_cg.py,sha256=YAEAD_8YU_H8Y-o6JI0Ywgk-kpAQOFBQm2Bjzaz9Bjs,2865
+torchzero/modules/second_order/nystrom.py,sha256=aM6dlDv7znGYNXZgKN6B6AhZ1Tpp01JMs83B1hcXE3w,6061
+torchzero/modules/smoothing/__init__.py,sha256=tUTGN0A-EQC7xuLV2AuHFWk-t7D6jIJlpV_3qyfRqLk,80
+torchzero/modules/smoothing/gaussian.py,sha256=KbCgRXGntdPbt4-ojalrHkniYgYXk2294b-2C4MIFi8,6109
+torchzero/modules/smoothing/laplacian.py,sha256=Vp2EnCQhyfGc3CbyOLc6_ZiVx_jvnOISf9vlHkIH4Jo,4998
+torchzero/modules/weight_decay/__init__.py,sha256=j2Vq3DDxLYIPJmXWgAJ6dL-rXzcDEZxxvhJqRT3H0-U,95
+torchzero/modules/weight_decay/weight_decay.py,sha256=UFL9W5w5nzTZGWvCwyGLe9UWBKN8FTClme1Klt7XZPw,3034
+torchzero/modules/wrappers/__init__.py,sha256=6b5Ac-8u18IVp_Jnw1T1xQExwpQhpQ0JwNV9GyC_Yj8,31
+torchzero/modules/wrappers/optim_wrapper.py,sha256=-wNI-fN8eaMSkvPIcPa34yxH0St5aLn7jaaLeh2DUsM,3569
+torchzero/optim/__init__.py,sha256=aXf7EkywqYiR50I4QeeVXro9aBhKiqfbY_BCia59sgU,46
+torchzero/optim/utility/__init__.py,sha256=pUacok4XmebfxofE-QWZLgViajsU-3JkXcWi9OS-Jrw,24
+torchzero/optim/utility/split.py,sha256=ZbazNuMTYunm75V_5ard0A_LletGaYAg-Pm2rANJKrE,1610
+torchzero/optim/wrappers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+torchzero/optim/wrappers/directsearch.py,sha256=Y2-7Sy4mYRPXPh0FTlsY_XOk5pCGjZsnbrlWCPZNp6A,10141
+torchzero/optim/wrappers/fcmaes.py,sha256=TQvIktXV8ldy6smBX-S7ZcQEbSmSZyj567TuYShbvJg,3513
+torchzero/optim/wrappers/mads.py,sha256=lC7edtrFS37PgmX7z9-eoqw6prl0k5BDB4NVBVQXJWE,2945
+torchzero/optim/wrappers/nevergrad.py,sha256=qslMb-4_kfjU3Dd0UbbzE2SdLViil3Qjo2v0FtPE3Fg,4000
+torchzero/optim/wrappers/nlopt.py,sha256=AaVEKfjbrt5DFION44_-g-jQAoVi4lCvBBPU5UDGO9Q,8151
+torchzero/optim/wrappers/optuna.py,sha256=YN1I3rzsi20A9963pWNWd7W75FkxalVb5z5fCRQeWA0,2280
+torchzero/optim/wrappers/scipy.py,sha256=pR26v8v0a-o2u0sbsKXpZ9JUKqXMaaI8gGLI8xYx3-s,19239
+torchzero/utils/__init__.py,sha256=7beAjXvnmBQoy5hwYHY_PBUtrrbYb9Z7-KrYgfcFkPE,844
+torchzero/utils/compile.py,sha256=N8AWLv_7oBUHYornmvvx_L4uynjiD-x5Hj1tBwei3-w,5127
+torchzero/utils/derivatives.py,sha256=sAVd0Q1xmIPpo_AxRuoow66Hy_3goX_9o3lQK_1TyW0,16909
+torchzero/utils/numberlist.py,sha256=cbG0UsSb9WCRxVhw8sd7Yf0bDy_gSqtghiJtkUxIO6U,6139
+torchzero/utils/ops.py,sha256=n4Su1sbgTzlHczuPEHkuWenTtNBCa_MvlQ_hCZkIPnQ,314
+torchzero/utils/optimizer.py,sha256=r52qu6pEcRH4lCXVlLxW5IweA6L-VrQj6RCMfdhzRpw,12466
+torchzero/utils/optuna_tools.py,sha256=F-1Xg0n_29MVEb6lqgUFFNIl9BNJ6MOdIJPduoNH4JU,1325
+torchzero/utils/params.py,sha256=nQo270aOURU7rJ_D102y2pSXbzhJPK0Z_ehx4mZBMes,5784
+torchzero/utils/python_tools.py,sha256=T5W7MpR7pNXiWSVw7gj-UuE9Ch0p9LRWuUZfg9Vtb-I,2794
+torchzero/utils/tensorlist.py,sha256=qSbiliVo1euFAksdHHHRbPUdYYxfkw1dvhpXj71wGy0,53162
+torchzero/utils/torch_tools.py,sha256=ohqnnZRlqdfp5PAfMSbQDIEKygW0_ARjxSEBp3Zo9nU,4756
+torchzero/utils/linalg/__init__.py,sha256=Dzbho3_z7JDdKzYD-QdLArg0ZEoC2BVGdlE3JoAnXHQ,272
+torchzero/utils/linalg/benchmark.py,sha256=wiIMn-GY2xxWbHVf8CPbJddUPeUPq9OUDkvbp1iILYI,479
+torchzero/utils/linalg/matrix_funcs.py,sha256=-LecWrPWbJvfeCgIzUhfWARa2aSZvJ12lHX7Jno38O4,3099
+torchzero/utils/linalg/orthogonalize.py,sha256=mDCkET7qgDZqf_y6oPYAK3d2L5HrB8gzOFPl0YoONaY,399
+torchzero/utils/linalg/qr.py,sha256=L-RXuYV-SIHI-Llq4y1rQ_Tz-yamds0_QNZeHapbjNE,2507
+torchzero/utils/linalg/solve.py,sha256=P0PMi0zro3G3Rd0X-JeoLk7tqYDB0js0aB4bpQ0OABU,5235
+torchzero/utils/linalg/svd.py,sha256=wBxl-JSciINV-N6zvM4SGdveqMr6idq51h68LyQQRYg,660
+torchzero-0.3.10.dist-info/licenses/LICENSE,sha256=r9ZciAoZoqKC_FNADE0ORukj1p1XhLXEbegdsAyqhJs,1087
+torchzero-0.3.10.dist-info/METADATA,sha256=_J7AbrIa-nD6UWbuydCwxAnSpKcC9O1Vp_rM896ZkYQ,14081
+torchzero-0.3.10.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+torchzero-0.3.10.dist-info/top_level.txt,sha256=YDdpIOb7HyKV9THOtOYsFFMTbxvCO0kiol4-83tDj-A,21
+torchzero-0.3.10.dist-info/RECORD,,

{torchzero-0.3.9.dist-info → torchzero-0.3.10.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.8.0)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

torchzero/core/preconditioner.py

@@ -1,138 +0,0 @@
-from abc import ABC, abstractmethod
-from collections import ChainMap, defaultdict
-from collections.abc import Mapping, Sequence
-from typing import Any, overload, final
-
-import torch
-
-from .module import Module, Chainable, Vars
-from .transform import apply, Transform, Target
-from ..utils import TensorList, vec_to_tensors
-
-class Preconditioner(Transform):
-    """Abstract class for a preconditioner."""
-    def __init__(
-        self,
-        defaults: dict | None,
-        uses_grad: bool,
-        concat_params: bool = False,
-        update_freq: int = 1,
-        scale_first: bool = False,
-        inner: Chainable | None = None,
-        target: Target = "update",
-    ):
-        if defaults is None: defaults = {}
-        defaults.update(dict(__update_freq=update_freq, __concat_params=concat_params, __scale_first=scale_first))
-        super().__init__(defaults, uses_grad=uses_grad, target=target)
-
-        if inner is not None:
-            self.set_child('inner', inner)
-
-    @abstractmethod
-    def update(self, tensors: list[torch.Tensor], params:list[torch.Tensor], grads:list[torch.Tensor] | None, states: list[dict[str, Any]], settings: Sequence[Mapping[str, Any]]):
-        """updates the preconditioner with `tensors`, any internal state should be stored using `keys`"""
-
-    @abstractmethod
-    def apply(self, tensors:list[torch.Tensor], params:list[torch.Tensor], grads:list[torch.Tensor] | None, states: list[dict[str, Any]], settings: Sequence[Mapping[str, Any]]) -> list[torch.Tensor]:
-        """applies preconditioner to `tensors`, any internal state should be stored using `keys`"""
-
-
-    def _tensor_wise_transform(self, tensors:list[torch.Tensor], params:list[torch.Tensor], grads:list[torch.Tensor] | None, vars:Vars) -> list[torch.Tensor]:
-        step = self.global_state.get('__step', 0)
-        states = [self.state[p] for p in params]
-        settings = [self.settings[p] for p in params]
-        global_settings = settings[0]
-        update_freq = global_settings['__update_freq']
-
-        scale_first = global_settings['__scale_first']
-        scale_factor = 1
-        if scale_first and step == 0:
-            # initial step size guess from pytorch LBFGS
-            scale_factor = 1 / TensorList(tensors).abs().global_sum().clip(min=1)
-            scale_factor = scale_factor.clip(min=torch.finfo(tensors[0].dtype).eps)
-
-        # update preconditioner
-        if step % update_freq == 0:
-            self.update(tensors=tensors, params=params, grads=grads, states=states, settings=settings)
-
-        # step with inner
-        if 'inner' in self.children:
-            tensors = apply(self.children['inner'], tensors=tensors, params=params, grads=grads, vars=vars)
-
-        # apply preconditioner
-        tensors = self.apply(tensors=tensors, params=params, grads=grads, states=states, settings=settings)
-
-        # scale initial step, when preconditioner might not have been applied
-        if scale_first and step == 0:
-            torch._foreach_mul_(tensors, scale_factor)
-
-        self.global_state['__step'] = step + 1
-        return tensors
-
-    def _concat_transform(self, tensors:list[torch.Tensor], params:list[torch.Tensor], grads:list[torch.Tensor] | None, vars:Vars) -> list[torch.Tensor]:
-        step = self.global_state.get('__step', 0)
-        tensors_vec = torch.cat([t.ravel() for t in tensors])
-        params_vec = torch.cat([p.ravel() for p in params])
-        grads_vec = [torch.cat([g.ravel() for g in grads])] if grads is not None else None
-
-        states = [self.state[params[0]]]
-        settings = [self.settings[params[0]]]
-        global_settings = settings[0]
-        update_freq = global_settings['__update_freq']
-
-        scale_first = global_settings['__scale_first']
-        scale_factor = 1
-        if scale_first and step == 0:
-            # initial step size guess from pytorch LBFGS
-            scale_factor = 1 / tensors_vec.abs().sum().clip(min=1)
-            scale_factor = scale_factor.clip(min=torch.finfo(tensors_vec.dtype).eps)
-
-        # update preconditioner
-        if step % update_freq == 0:
-            self.update(tensors=[tensors_vec], params=[params_vec], grads=grads_vec, states=states, settings=settings)
-
-        # step with inner
-        if 'inner' in self.children:
-            tensors = apply(self.children['inner'], tensors=tensors, params=params, grads=grads, vars=vars)
-            tensors_vec = torch.cat([t.ravel() for t in tensors]) # have to recat
-
-        # apply preconditioner
-        tensors_vec = self.apply(tensors=[tensors_vec], params=[params_vec], grads=grads_vec, states=states, settings=settings)[0]
-
-        # scale initial step, when preconditioner might not have been applied
-        if scale_first and step == 0:
-            tensors_vec *= scale_factor
-
-        tensors = vec_to_tensors(vec=tensors_vec, reference=tensors)
-        self.global_state['__step'] = step + 1
-        return tensors
-
-    @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        concat_params = self.settings[params[0]]['__concat_params']
-        if concat_params: return self._concat_transform(tensors, params, grads, vars)
-        return self._tensor_wise_transform(tensors, params, grads, vars)
-
-class TensorwisePreconditioner(Preconditioner, ABC):
-    @abstractmethod
-    def update_tensor(self, tensor: torch.Tensor, param:torch.Tensor, grad: torch.Tensor | None, state: dict[str, Any], settings: Mapping[str, Any]):
-        """update preconditioner with `tensor`"""
-
-    @abstractmethod
-    def apply_tensor(self, tensor: torch.Tensor, param:torch.Tensor, grad: torch.Tensor | None, state: dict[str, Any], settings: Mapping[str, Any]) -> torch.Tensor:
-        """apply preconditioner to `tensor`"""
-
-    @final
-    def update(self, tensors, params, grads, states, settings):
-        if grads is None: grads = [None]*len(tensors)
-        for t,p,g,state,setting in zip(tensors, params, grads, states, settings):
-            self.update_tensor(t, p, g, state, setting)
-
-    @final
-    def apply(self, tensors, params, grads, states, settings):
-        preconditioned = []
-        if grads is None: grads = [None]*len(tensors)
-        for t,p,g,state,setting in zip(tensors, params, grads, states, settings):
-            preconditioned.append(self.apply_tensor(t, p, g, state, setting))
-        return preconditioned
-