torchzero 0.3.15__py3-none-any.whl → 0.4.0__py3-none-any.whl

torchzero/core/modular.py

@@ -1,27 +1,16 @@
 
 import warnings
-from abc import ABC, abstractmethod
-from collections import ChainMap, defaultdict
-from collections.abc import Callable, Iterable, MutableMapping, Sequence
-from operator import itemgetter
-from typing import TYPE_CHECKING, Any, Literal, cast, final, overload
+from collections import ChainMap
+from collections.abc import MutableMapping
+from typing import Any
 
 import torch
 
-from ..utils import (
-    Init,
-    ListLike,
-    Params,
-    _make_param_groups,
-    get_state_vals,
-    vec_to_tensors,
-)
-from ..utils.derivatives import flatten_jacobian, hvp, hvp_fd_central, hvp_fd_forward
-from ..utils.linalg.linear_operator import LinearOperator
-from ..utils.python_tools import flatten
-from .module import Chainable, Module
-from .var import Var
+from ..utils.params import Params, _make_param_groups
 from .functional import step
+from .module import Chainable, Module
+from .objective import Objective
+
 
 class _EvalCounterClosure:
     """keeps track of how many times closure has been evaluated, and sets closure return"""
@@ -32,7 +21,7 @@ class _EvalCounterClosure:
 
     def __call__(self, *args, **kwargs):
         if self.closure is None:
-            raise RuntimeError("One of the modules requires closure to be passed to the step method")
+            raise RuntimeError("closure is None in _EvalCounterClosure, and this can't happen")
 
         v = self.closure(*args, **kwargs)
 
@@ -44,17 +33,17 @@ class _EvalCounterClosure:
         return v
 
 
-def unroll_modules(*modules: Chainable) -> list[Module]:
-    unrolled = []
+def flatten_modules(*modules: Chainable) -> list[Module]:
+    flat = []
 
     for m in modules:
         if isinstance(m, Module):
-            unrolled.append(m)
-            unrolled.extend(unroll_modules(list(m.children.values())))
+            flat.append(m)
+            flat.extend(flatten_modules(list(m.children.values())))
         else:
-            unrolled.extend(unroll_modules(*m))
+            flat.extend(flatten_modules(*m))
 
-    return unrolled
+    return flat
 
 
 # have to inherit from Modular to support lr schedulers
@@ -83,7 +72,7 @@ class Modular(torch.optim.Optimizer):
         self.modules = modules
         """Top-level modules providedduring initialization."""
 
-        self.unrolled_modules = unroll_modules(self.modules)
+        self.flat_modules = flatten_modules(self.modules)
         """A flattened list of all modules including all children."""
 
         param_groups = _make_param_groups(params, differentiable=False)
@@ -92,7 +81,7 @@ class Modular(torch.optim.Optimizer):
         Each element in the list is ChainDict's 2nd map of a module."""
 
         # make sure there is no more than a single learning rate module
-        lr_modules = [m for m in self.unrolled_modules if 'lr' in m.defaults]
+        lr_modules = [m for m in self.flat_modules if 'lr' in m.defaults]
         if len(lr_modules) > 1:
             warnings.warn(f'multiple learning rate modules detected: {lr_modules}. This may lead to componding of learning rate multiplication with per-parameter learning rates and schedulers.')
 
@@ -100,13 +89,13 @@ class Modular(torch.optim.Optimizer):
         for group in param_groups:
             for k in group:
                 if k in ('params', 'lr'): continue
-                modules_with_k = [m for m in self.unrolled_modules if k in m.defaults and k not in m._overridden_keys]
+                modules_with_k = [m for m in self.flat_modules if k in m.defaults and k not in m._overridden_keys]
                 if len(modules_with_k) > 1:
                     warnings.warn(f'`params` has a `{k}` key, and multiple modules have that key: {modules_with_k}. If you intended to only set `{k}` to one of them, use `module.set_param_groups(params)`')
 
         # defaults for schedulers
         defaults = {}
-        for m in self.unrolled_modules: defaults.update(m.defaults)
+        for m in self.flat_modules: defaults.update(m.defaults)
         super().__init__(param_groups, defaults=defaults)
 
         # note - this is what super().__init__(param_groups, defaults=defaults) does:
@@ -146,7 +135,7 @@ class Modular(torch.optim.Optimizer):
 
         for p in proc_param_group['params']:
             # updates global per-parameter setting overrides (medium priority)
-            self._per_parameter_global_settings[p] = [m.settings[p].maps[1] for m in self.unrolled_modules]
+            self._per_parameter_global_settings[p] = [m.settings[p].maps[1] for m in self.flat_modules]
 
     def state_dict(self):
         all_params = [p for g in self.param_groups for p in g['params']]
@@ -163,7 +152,7 @@ class Modular(torch.optim.Optimizer):
             "params": all_params,
             "groups": groups,
             "defaults": self.defaults,
-            "modules": {i: m.state_dict() for i, m in enumerate(self.unrolled_modules)}
+            "modules": {i: m.state_dict() for i, m in enumerate(self.flat_modules)}
         }
         return state_dict
 
@@ -183,7 +172,7 @@ class Modular(torch.optim.Optimizer):
             self.add_param_group(group)
 
         id_to_tensor = {state_dict['idx_to_id'][i]: p for i,p in enumerate(state_dict['params'])}
-        for m, sd in zip(self.unrolled_modules, state_dict['modules'].values()):
+        for m, sd in zip(self.flat_modules, state_dict['modules'].values()):
             m._load_state_dict(sd, id_to_tensor)
 
 
@@ -201,35 +190,42 @@ class Modular(torch.optim.Optimizer):
                 if not p.requires_grad: continue
                 for map in self._per_parameter_global_settings[p]: map.update(settings)
 
-        # create var
+        # create Objective
         params = [p for g in self.param_groups for p in g['params'] if p.requires_grad]
-        var = Var(params=params, closure=_EvalCounterClosure(self, closure), model=self.model, current_step=self.current_step, modular=self, loss=loss, storage=kwargs)
 
-        # if closure is None, assume backward has been called and gather grads
-        if closure is None:
-            var.grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
-            self.num_evaluations += 1
+        counter_closure = None
+        if closure is not None:
+            counter_closure = _EvalCounterClosure(self, closure)
 
-        if len(self.modules) == 0: raise RuntimeError("There are no modules in this `Modular` optimizer")
+        objective = Objective(
+            params=params, closure=counter_closure, model=self.model,
+            current_step=self.current_step, modular=self, loss=loss, storage=kwargs
+        )
 
-        # step
-        var = step(var, self.modules)
+        # step with all modules
+        objective = step(objective, self.modules)
 
-        # apply update
-        if not var.skip_update:
-            with torch.no_grad():
-                torch._foreach_sub_(params, var.get_update())
+        # apply update to parameters unless `objective.skip_update = True`
+        # this does:
+        # if not objective.skip_update:
+        #   torch._foreach_sub_(objective.params, objective.get_updates())
+        objective.update_parameters()
 
         # update attributes
-        self.attrs.update(var.attrs)
-        if var.should_terminate is not None: self.should_terminate = var.should_terminate
-
-        # hooks
-        for hook in var.post_step_hooks:
-            hook(self, var)
+        self.attrs.update(objective.attrs)
+        if objective.should_terminate is not None:
+            self.should_terminate = objective.should_terminate
 
         self.current_step += 1
-        #return var.loss if var.loss is not None else var.loss_approx
+
+        # apply hooks
+        # this does:
+        # for hook in objective.post_step_hooks:
+        #     hook(objective, modules)
+        objective.apply_post_step_hooks(self.modules)
+
+        # return the first closure evaluation return
+        # could return loss if it was passed but that's pointless
         return self._closure_return
 
     def __repr__(self):

torchzero/core/module.py

@@ -1,24 +1,18 @@
 import warnings
 from abc import ABC, abstractmethod
 from collections import ChainMap, defaultdict
-from collections.abc import Callable, Iterable, MutableMapping, Sequence
-from operator import itemgetter
-from typing import Any, Literal, cast, final, overload
+from collections.abc import Callable, Iterable, Sequence
+from typing import Any, overload, TYPE_CHECKING
 
 import torch
 
-from ..utils import (
-    Init,
-    ListLike,
-    Params,
-    _make_param_groups,
-    get_state_vals,
-    vec_to_tensors,
-)
-from ..utils.derivatives import flatten_jacobian, hvp, hvp_fd_central, hvp_fd_forward
-from ..utils.linalg.linear_operator import LinearOperator
-from ..utils.python_tools import flatten
-from .var import Var
+from ..linalg.linear_operator import LinearOperator
+from ..utils.optimizer import Init, ListLike, get_state_vals
+from ..utils.params import Params, _make_param_groups
+from .functional import step_tensors
+
+if TYPE_CHECKING:
+    from .objective import Objective
 
 
 class Module(ABC):
@@ -36,6 +30,7 @@ class Module(ABC):
 """
     def __init__(self, defaults: dict[str, Any] | None = None):
         if defaults is None: defaults = {}
+        if any(isinstance(v, Module) for v in defaults.values()): raise RuntimeError("Passed a module to defaults")
         self.defaults: dict[str, Any] = defaults
 
         # settings are stored like state in per-tensor defaultdict, with per-parameter overrides possible
@@ -55,7 +50,7 @@ class Module(ABC):
         """A dictionary of child modules."""
 
         self._overridden_keys = set()
-        """tracks keys overridden with `set_param_groups`, only used to not give a warning"""
+        """tracks keys overridden with ``set_param_groups``, only used to not give a warning"""
 
 
     def set_param_groups(self, param_groups: Params):
@@ -71,7 +66,12 @@ class Module(ABC):
                 self.settings[param].maps[0].update(settings) # set module-specific per-parameter settings
         return self
 
-    def set_child(self, key: str, module: "Module | Sequence[Module]"):
+    def set_child(self, key: str, module: "Module | Sequence[Module] | None"):
+        if key in self.children:
+            warnings.warn(f"set_child overwriting child `{key}`")
+
+        if module is None: return
+
         from .chain import maybe_chain
         self.children[key] = maybe_chain(module)
 
@@ -85,6 +85,62 @@ class Module(ABC):
     def get_children_sequence(self, prefix = 'module_'):
         return [self.children[f'{prefix}{i}'] for i in range(len(self.children)) if f'{prefix}{i}' in self.children]
 
+    def inner_step(
+        self,
+        key: str,
+        objective: "Objective",
+        must_exist: bool = True,
+    ) -> "Objective":
+        """Passes ``objective`` to child and returns it."""
+        child = self.children.get(key, None)
+
+        if child is None:
+            if must_exist: raise KeyError(f"child `{key}` doesn't exist")
+            return objective
+
+        return child.step(objective)
+
+
+    def inner_step_tensors(
+        self,
+        key: str,
+        tensors: list[torch.Tensor],
+        clone: bool,
+        params: Iterable[torch.Tensor] | None = None,
+        grads: Sequence[torch.Tensor] | None = None,
+        loss: torch.Tensor | None = None,
+        closure: Callable | None = None,
+        objective: "Objective | None" = None,
+        must_exist: bool = True
+    ) -> list[torch.Tensor]:
+        """Steps with child module. Can be used to apply transforms to any internal buffers.
+
+        If ``objective`` is specified, other attributes shouldn't to be specified.
+
+        Args:
+            key (str): Child module key.
+            tensors (Sequence[torch.Tensor]): tensors to pass to child module.
+            clone (bool):
+                If ``key`` exists, whether to clone ``tensors`` to avoid modifying buffers in-place.
+                If ``key`` doesn't exist, ``tensors`` are always returned without cloning
+            params (Iterable[torch.Tensor] | None, optional): pass None if ``tensors`` have different shape. Defaults to None.
+            grads (Sequence[torch.Tensor] | None, optional): grads. Defaults to None.
+            loss (torch.Tensor | None, optional): loss. Defaults to None.
+            closure (Callable | None, optional): closure. Defaults to None.
+            must_exist (bool, optional): if True, if ``key`` doesn't exist, raises ``KeyError``. Defaults to True.
+        """
+
+        child = self.children.get(key, None)
+
+        if child is None:
+            if must_exist: raise KeyError(f"child `{key}` doesn't exist")
+            return tensors
+
+        if clone: tensors = [t.clone() for t in tensors]
+        return step_tensors(modules=child, tensors=tensors, params=params, grads=grads,
+                            loss=loss, closure=closure, objective=objective)
+
+
     def __repr__(self):
         s = self.__class__.__name__
         if self.children:
@@ -106,7 +162,6 @@ class Module(ABC):
 
     def get_settings(self, params: Sequence[torch.Tensor], key: str | list[str] | tuple[str,...], key2: str | None = None,
                      *keys: str, cls: type[ListLike] = list) -> ListLike | list[ListLike]:
-        # if isinstance(params, Vars): params = params.params
         return get_state_vals(self.settings, params, key, key2, *keys, must_exist=True, cls=cls) # pyright:ignore[reportArgumentType]
 
 
@@ -176,13 +231,8 @@ class Module(ABC):
             - if state_keys has multiple keys and keys has a single key, return cls.
             - if state_keys has multiple keys and keys has multiple keys, return list of cls.
         """
-        # if isinstance(params, Vars): params = params.params
         return get_state_vals(self.state, params, key, key2, *keys, must_exist=must_exist, init=init, cls=cls) # pyright:ignore[reportArgumentType]
 
-    # def first_setting(self, *keys:str, params:Sequence[torch.Tensor]):
-    #     # if isinstance(params, Vars): params = params.params
-    #     return itemgetter(*keys)(self.settings[params[0]])
-
     def clear_state_keys(self, *keys:str):
         for s in self.state.values():
             for k in keys:
@@ -248,36 +298,73 @@ class Module(ABC):
         # extra info
         self._extra_unpack(state_dict['extra'])
 
-    # ---------------------------- OVERRIDABLE METHODS --------------------------- #
-    def step(self, var: Var) -> Var:
-        """performs a step, returns new ``var`` but may update it in-place."""
-        self.update(var)
-        return self.apply(var)
+    def get_generator(self, device: torch.types.Device, seed: int | None):
+        """If ``seed=None``, returns ``None``.
+
+        Otherwise, if generator on this device and with this seed hasn't been created,
+        creates it and stores in global state.
+
+        Returns ``torch.Generator``."""
+        if seed is None: return None
 
-    def update(self, var:Var) -> Any:
-        """Updates the internal state of this module. This should not modify ``var.update``.
+        if device is None: device_obj = torch.get_default_device()
+        else: device_obj = torch.device(device)
+        key = f"__generator-{seed}-{device_obj.type}:{device_obj.index}"
+
+        if key not in self.global_state:
+            self.global_state[key] = torch.Generator(device).manual_seed(seed)
+
+        return self.global_state[key]
+
+    def increment_counter(self, key: str, start: int):
+        """first value is ``start``"""
+        value = self.global_state.get(key, start - 1) + 1
+        self.global_state[key] = value
+        return value
+
+    # ---------------------------- OVERRIDABLE METHODS --------------------------- #
+    def update(self, objective:"Objective") -> None:
+        """Updates internal state of this module. This should not modify ``objective.update``.
 
         Specifying ``update`` and ``apply`` methods is optional and allows certain meta-modules to be used,
-        such as ``tz.m.Online`` or trust regions. Alternatively, simply override the ``step`` method.
+        such as ``tz.m.Online`` or trust regions. Alternatively, define all logic within the ``apply`` method.
+
+        ``update`` is guaranteed to be called at least once before ``apply``.
+
+        Args:
+            objective (Objective): ``Objective`` object
         """
 
-    def apply(self, var: Var) -> Var:
-        """Applies this module to ``var.get_update()``.
-        This should not modify the internal state of this module if possible.
+    @abstractmethod
+    def apply(self, objective: "Objective") -> "Objective":
+        """Updates ``objective`` using the internal state of this module.
+
+        If ``update`` method is defined, ``apply`` shouldn't modify the internal state of this module if possible.
 
         Specifying ``update`` and ``apply`` methods is optional and allows certain meta-modules to be used,
-        such as ``tz.m.Online`` or trust regions. Alternatively, simply override the ``step`` method.
+        such as ``tz.m.Online`` or trust regions. Alternatively, define all logic within the ``apply`` method.
+
+        ``update`` is guaranteed to be called at least once before ``apply``.
+
+        Args:
+            objective (Objective): ``Objective`` object
         """
-        return self.step(var)
+        # if apply is empty, it should be defined explicitly.
+        raise NotImplementedError(f"{self.__class__.__name__} doesn't implement `apply`.")
+
+    def step(self, objective: "Objective") -> "Objective":
+        """Perform a step with this module. Calls ``update``, then ``apply``."""
+        self.update(objective)
+        return self.apply(objective)
 
-    def get_H(self, var: Var) -> LinearOperator | None:
+    def get_H(self, objective: "Objective") -> LinearOperator | None:
         """returns a ``LinearOperator`` corresponding to hessian or hessian approximation.
         The hessian approximation is assumed to be for all parameters concatenated to a vector."""
         # if this method is not defined it searches in children
         # this should be overwritten to return None if child params are different from this modules params
         H = None
         for k,v in self.children.items():
-            H_v = v.get_H(var)
+            H_v = v.get_H(objective)
 
             if (H is not None) and (H_v is not None):
                 raise RuntimeError(f"Two children of {self} have a hessian, second one is {k}={v}")
@@ -307,21 +394,14 @@ class Module(ABC):
         """
         for c in self.children.values(): c.reset_for_online()
 
-    def _extra_pack(self):
-        """extra information to store in state_dict of this optimizer.
-        Will be passed to ``_extra_unpack`` when loading the state_dict."""
+    def _extra_pack(self) -> dict:
+        """extra information to store in ``state_dict`` of this optimizer.
+        Will be passed to ``_extra_unpack`` when loading the ``state_dict``."""
         return {}
 
-    def _extra_unpack(self, x):
-        """``_extra_pack`` return will be passed to this method when loading state_dict.
+    def _extra_unpack(self, d: dict):
+        """``_extra_pack`` return will be passed to this method when loading ``state_dict``.
         This method is called after loading the rest of the state dict"""
 
-    def get_generator(self, device: torch.types.Device, seed: int | None):
-        if seed is None: return None
-
-        if 'generator' not in self.global_state:
-            self.global_state['generator'] = torch.Generator(device).manual_seed(seed)
-
-        return self.global_state['generator']
 
 Chainable = Module | Sequence[Module]