lucid-dl 2.11.3__py3-none-any.whl → 2.11.4__py3-none-any.whl

lucid/nn/module.py

@@ -13,7 +13,22 @@ from typing import (
 from collections import OrderedDict
 
 from lucid._tensor import Tensor
-from lucid.types import _ArrayOrScalar, _NumPyArray, _DeviceType
+from lucid.types import (
+    _ArrayOrScalar,
+    _BackwardHook,
+    _DeviceType,
+    _ForwardHook,
+    _ForwardHookKwargs,
+    _ForwardPreHook,
+    _ForwardPreHookKwargs,
+    _FullBackwardHook,
+    _FullBackwardPreHook,
+    _LoadStateDictPostHook,
+    _LoadStateDictPreHook,
+    _NumPyArray,
+    _StateDictHook,
+    _StateDictPreHook,
+)
 
 import lucid.nn as nn
 
@@ -30,26 +45,6 @@ __all__ = [
 ]
 
 
-_ForwardPreHook = Callable[["Module", tuple[Any, ...]], tuple[Any, ...] | None]
-_ForwardPreHookKwargs = Callable[
-    ["Module", tuple[Any, ...], dict[str, Any]],
-    tuple[tuple[Any, ...], dict[str, Any]] | None,
-]
-_ForwardHook = Callable[["Module", tuple[Any, ...], Any], Any | None]
-_ForwardHookKwargs = Callable[
-    ["Module", tuple[Any, ...], dict[str, Any], Any], Any | None
-]
-
-_BackwardHook = Callable[[Tensor, _NumPyArray], None]
-_FullBackwardPreHook = Callable[
-    ["Module", tuple[_NumPyArray | None, ...]], tuple[_NumPyArray | None, ...] | None
-]
-_FullBackwardHook = Callable[
-    ["Module", tuple[_NumPyArray | None, ...], tuple[_NumPyArray | None, ...]],
-    tuple[_NumPyArray | None, ...] | None,
-]
-
-
 class Module:
     _registry_map: dict[Type, OrderedDict[str, Any]] = {}
     _alt_name: str = ""
@@ -70,10 +65,17 @@ class Module:
             tuple[_ForwardPreHook | _ForwardPreHookKwargs, bool]
         ] = []
         self._forward_hooks: list[tuple[_ForwardHook | _ForwardHookKwargs, bool]] = []
+
         self._backward_hooks: list[_BackwardHook] = []
         self._full_backward_pre_hooks: list[_FullBackwardPreHook] = []
         self._full_backward_hooks: list[_FullBackwardHook] = []
 
+        self._state_dict_pre_hooks: list[_StateDictPreHook] = []
+        self._state_dict_hooks: list[_StateDictHook] = []
+
+        self._load_state_dict_pre_hooks: list[_LoadStateDictPreHook] = []
+        self._load_state_dict_post_hooks: list[_LoadStateDictPostHook] = []
+
         self._state_dict_pass_attr = set()
 
     def __setattr__(self, name: str, value: Any) -> None:
@@ -155,6 +157,26 @@ class Module:
         self._full_backward_hooks.append(hook)
         return lambda: self._full_backward_hooks.remove(hook)
 
+    def register_state_dict_pre_hook(self, hook: _StateDictPreHook) -> Callable:
+        self._state_dict_pre_hooks.append(hook)
+        return lambda: self._state_dict_pre_hooks.remove(hook)
+
+    def register_state_dict_hook(self, hook: _StateDictHook) -> Callable:
+        self._state_dict_hooks.append(hook)
+        return lambda: self._state_dict_hooks.remove(hook)
+
+    def register_load_state_dict_pre_hook(
+        self, hook: _LoadStateDictPreHook
+    ) -> Callable:
+        self._load_state_dict_pre_hooks.append(hook)
+        return lambda: self._load_state_dict_pre_hooks.remove(hook)
+
+    def register_load_state_dict_post_hook(
+        self, hook: _LoadStateDictPostHook
+    ) -> Callable:
+        self._load_state_dict_post_hooks.append(hook)
+        return lambda: self._load_state_dict_post_hooks.remove(hook)
+
     def reset_parameters(self) -> None:
         for param in self.parameters():
             param.zero()
@@ -231,6 +253,9 @@ class Module:
         prefix: str = "",
         keep_vars: bool = False,
     ) -> OrderedDict:
+        for hook in self._state_dict_pre_hooks:
+            hook(self, prefix, keep_vars)
+
         if destination is None:
             destination = OrderedDict()
 
@@ -249,9 +274,15 @@ class Module:
             if key in self._state_dict_pass_attr:
                 del destination[key]
 
+        for hook in self._state_dict_hooks:
+            hook(self, destination, prefix, keep_vars)
+
         return destination
 
     def load_state_dict(self, state_dict: OrderedDict, strict: bool = True) -> None:
+        for hook in self._load_state_dict_pre_hooks:
+            hook(self, state_dict, strict)
+
         own_state = self.state_dict(keep_vars=True)
 
         missing_keys = set(own_state.keys()) - set(state_dict.keys())
@@ -277,6 +308,9 @@ class Module:
             elif strict:
                 raise KeyError(f"Unexpected key '{key}' in state_dict.")
 
+        for hook in self._load_state_dict_post_hooks:
+            hook(self, missing_keys, unexpected_keys, strict)
+
     def __call__(self, *args: Any, **kwargs: Any) -> Tensor | tuple[Tensor, ...]:
         for hook, with_kwargs in self._forward_pre_hooks:
             if with_kwargs:

lucid/types.py

@@ -6,8 +6,10 @@ from typing import (
     Sequence,
     Literal,
     TypeAlias,
+    TYPE_CHECKING,
     runtime_checkable,
 )
+from collections import OrderedDict
 import re
 
 import numpy as np
@@ -76,6 +78,62 @@ class _TensorLike(Protocol):
     ) -> None: ...
 
 
+@runtime_checkable
+class _ModuleHookable(Protocol):
+    def register_forward_pre_hook(
+        self, hook: Callable, *, with_kwargs: bool = False
+    ) -> Callable: ...
+
+    def register_forward_hook(
+        self, hook: Callable, *, with_kwargs: bool = False
+    ) -> Callable: ...
+
+    def register_backward_hook(self, hook: Callable) -> Callable: ...
+
+    def register_full_backward_pre_hook(self, hook: Callable) -> Callable: ...
+
+    def register_full_backward_hook(self, hook: Callable) -> Callable: ...
+
+    def register_state_dict_pre_hook(self, hook: Callable) -> Callable: ...
+
+    def register_state_dict_hook(self, hook: Callable) -> Callable: ...
+
+    def register_load_state_dict_pre_hook(self, hook: Callable) -> Callable: ...
+
+    def register_load_state_dict_post_hook(self, hook: Callable) -> Callable: ...
+
+
+_ForwardPreHook: TypeAlias = Callable[
+    [_ModuleHookable, tuple[Any, ...]], tuple[Any, ...] | None
+]
+_ForwardPreHookKwargs: TypeAlias = Callable[
+    [_ModuleHookable, tuple[Any, ...], dict[str, Any]],
+    tuple[tuple[Any, ...], dict[str, Any]] | None,
+]
+_ForwardHook: TypeAlias = Callable[[_ModuleHookable, tuple[Any, ...], Any], Any | None]
+_ForwardHookKwargs: TypeAlias = Callable[
+    [_ModuleHookable, tuple[Any, ...], dict[str, Any], Any], Any | None
+]
+
+_BackwardHook: TypeAlias = Callable[[_TensorLike, _NumPyArray], None]
+_FullBackwardPreHook: TypeAlias = Callable[
+    [_ModuleHookable, tuple[_NumPyArray | None, ...]],
+    tuple[_NumPyArray | None, ...] | None,
+]
+_FullBackwardHook: TypeAlias = Callable[
+    [_ModuleHookable, tuple[_NumPyArray | None, ...], tuple[_NumPyArray | None, ...]],
+    tuple[_NumPyArray | None, ...] | None,
+]
+
+_StateDictPreHook: TypeAlias = Callable[[_ModuleHookable, str, bool], None]
+_StateDictHook: TypeAlias = Callable[[_ModuleHookable, OrderedDict, str, bool], None]
+
+_LoadStateDictPreHook: TypeAlias = Callable[[_ModuleHookable, OrderedDict, bool], None]
+_LoadStateDictPostHook: TypeAlias = Callable[
+    [_ModuleHookable, set[str], set[str], bool], None
+]
+
+
 class Numeric:
     def __init__(
         self, base_dtype: type[int | float | complex], bits: int | None

{lucid_dl-2.11.3.dist-info → lucid_dl-2.11.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lucid-dl
-Version: 2.11.3
+Version: 2.11.4
 Summary: Lumerico's Comprehensive Interface for Deep Learning
 Home-page: https://github.com/ChanLumerico/lucid
 Author: ChanLumerico
@@ -52,15 +52,31 @@ Whether you're a student, educator, or an advanced researcher seeking to demysti
 
   ```python
   def register_forward_pre_hook(self, hook: Callable, *, with_kwargs: bool = False)
-
+  ```
+  ```python
   def register_forward_hook(self, hook: Callable, *, with_kwargs: bool = False)
-
+  ```
+  ```python
   def register_backward_hook(self, hook: Callable)
-
+  ```
+  ```python
   def register_full_backward_pre_hook(self, hook: Callable)
-
+  ```
+  ```python
   def register_full_backward_hook(self, hook: Callable)
   ```
+  ```python
+  def register_state_dict_pre_hook(self, hook: Callable)
+  ```
+  ```python
+  def register_state_dict_hook(self, hook: Callable)
+  ```
+  ```python
+  def register_load_state_dict_pre_hook(self, hook: Callable)
+  ```
+  ```python
+  def register_load_state_dict_post_hook(self, hook: Callable)
+  ```
 
 ## 🔧 How to Install
 

{lucid_dl-2.11.3.dist-info → lucid_dl-2.11.4.dist-info}/RECORD

@@ -1,7 +1,7 @@
 lucid/__init__.py,sha256=EwNZkKALNS54RQw_D_1BtqbBo3kFMi4uSveSs6i-bdM,9161
 lucid/error.py,sha256=qnTiVuZm3c5-DIt-OOyobZ7RUm7E1K4NR0j998LG1ug,709
 lucid/port.py,sha256=Kt1YaSWef_eKF4KRj-UFhirvFC5urEESfYQ_BSlBZGE,3811
-lucid/types.py,sha256=XFeLUXhtMXGh9liRPB3X49PbJSK3bcOA6KN1gvCjksA,4818
+lucid/types.py,sha256=Zdz2r4ledouEG-6Gi6yEza5vSLyyTzZJn7AcRKbxy8o,6906
 lucid/_backend/__init__.py,sha256=n1bnYdeb_bNDBKASWGywTRa0Ne9hMAkal3AuVZJgovI,5
 lucid/_backend/core.py,sha256=neZF9uQlwNp-yHeyi0IbKlN556O-lsqThsIFExsf8-Y,11791
 lucid/_backend/metal.py,sha256=vQegTENuPjeAM_EIXfuOpnIXZBeMbGVxCFvm4s-_NNo,4215
@@ -78,7 +78,7 @@ lucid/models/seq2seq/__init__.py,sha256=wjsrhj4H_AcqwwbebAN8b68QBA8L6p1_12dkG299
 lucid/models/seq2seq/transformer.py,sha256=y5rerCs1s6jXTsVvbgscWScKpQKuSu1fezsBe7PNTRA,3513
 lucid/nn/__init__.py,sha256=_hk6KltQIJuWXowXstMSu3TjiaTP8zMLNvGpjnA9Mpw,182
 lucid/nn/fused.py,sha256=75fcXuo6fHSO-JtjuKhowhHSDr4qc5871WR63sUzH0g,5492
-lucid/nn/module.py,sha256=B2esFTdKJFjqsNUsFhVx_7IE-5nyJwW4b2I8QG1iLUk,25791
+lucid/nn/module.py,sha256=_EWtGkAuWWCPZ5f3t5pJOOzpi14gQBpP7JW2S8o4_GE,26855
 lucid/nn/parameter.py,sha256=NQS65YKn2B59wZbZIoT1mpDsU_F08y3yLi7hmV1B6yo,1232
 lucid/nn/util.py,sha256=Yw1iBSPrGV_r_F51qpqLYdafNE_hyaA0DPWYP-rjaig,1699
 lucid/nn/_kernel/__init__.py,sha256=n1bnYdeb_bNDBKASWGywTRa0Ne9hMAkal3AuVZJgovI,5
@@ -135,8 +135,8 @@ lucid/visual/graph.py,sha256=ZSlrJI3dQwYjz8XbgAfNd8-8YuH9Ji7Mz1J6UsnHTaI,4711
 lucid/visual/mermaid.py,sha256=87hFe4l9EYP6Cg2l2hP2INQiBHKkgVClH5nBWFY9ddY,26499
 lucid/weights/__init__.py,sha256=z1AikA3rOEeckWGkYWlcZkxNlJo9Xwa39PL6ly3hWnc,8801
 lucid/weights/__init__.pyi,sha256=lFonYC3cUx2Idolf3AEPnjFcyqcn3UDU84oJlZafqLY,3013
-lucid_dl-2.11.3.dist-info/licenses/LICENSE,sha256=vxRFYnVD1IeYtsvw-KmoElfqrjxKHv1h9YTvsG54loQ,1065
-lucid_dl-2.11.3.dist-info/METADATA,sha256=hffKVg1_fBZA5K3NcRn9KUkIv7B7ilHgr_9OJscZlR0,11898
-lucid_dl-2.11.3.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-lucid_dl-2.11.3.dist-info/top_level.txt,sha256=uzP_qBx9iNWIHKJRlElYcBLYVqMpdm9Q1Ma63QPYbFc,6
-lucid_dl-2.11.3.dist-info/RECORD,,
+lucid_dl-2.11.4.dist-info/licenses/LICENSE,sha256=vxRFYnVD1IeYtsvw-KmoElfqrjxKHv1h9YTvsG54loQ,1065
+lucid_dl-2.11.4.dist-info/METADATA,sha256=F8r0MrpLAlRuT0IK0RFLQCurlxcj3gUYLK2-tyKhAOI,12273
+lucid_dl-2.11.4.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+lucid_dl-2.11.4.dist-info/top_level.txt,sha256=uzP_qBx9iNWIHKJRlElYcBLYVqMpdm9Q1Ma63QPYbFc,6
+lucid_dl-2.11.4.dist-info/RECORD,,