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

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,11 +30,12 @@ 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
         # 0 - this module specific per-parameter setting overrides set via `set_param_groups` - highest priority
-        # 1 - global per-parameter setting overrides in param_groups passed to Modular - medium priority
+        # 1 - global per-parameter setting overrides in param_groups passed to Optimizer - medium priority
         # 2 - `defaults` - lowest priority
         self.settings: defaultdict[torch.Tensor, ChainMap[str, Any]] = defaultdict(lambda: ChainMap({}, {}, self.defaults))
         """per-parameter settings."""
@@ -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:
@@ -223,7 +273,7 @@ class Module(ABC):
         return state_dict
 
     def _load_state_dict(self, state_dict: dict[str, Any], id_to_tensor: dict[int, torch.Tensor]):
-        """loads state_dict, ``id_to_tensor`` is passed by ``Modular``"""
+        """loads state_dict, ``id_to_tensor`` is passed by ``Optimizer``"""
         # load state
         state = state_dict['state']
         self.state.clear()
@@ -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]