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

lucid/models/imgclf/crossvit.py

@@ -79,7 +79,7 @@ class _PatchEmbed(nn.Module):
                 f"Input image size {(H, W)} does not match with {self.img_size}."
             )
 
-        x = self.proj(x).flatten(axis=2).swapaxes(1, 2)
+        x = self.proj(x).flatten(start_axis=2).swapaxes(1, 2)
         return x
 
 

lucid/models/imgclf/efficientformer.py

@@ -80,7 +80,7 @@ class _Attention(nn.Module):
         y, x = lucid.meshgrid(
             lucid.arange(resolution[0]), lucid.arange(resolution[1]), indexing="ij"
         )
-        pos = lucid.stack([y, x]).flatten(axis=1)
+        pos = lucid.stack([y, x]).flatten(start_axis=1)
         rel_pos = lucid.abs(pos[..., :, None] - pos[..., None, :])
         rel_pos = (rel_pos[0] * resolution[1]) + rel_pos[1]
 
@@ -159,7 +159,7 @@ class _Downsample(nn.Module):
 
 class _Flatten(nn.Module):
     def forward(self, x: Tensor) -> Tensor:
-        x = x.flatten(axis=2).swapaxes(1, 2)
+        x = x.flatten(start_axis=2).swapaxes(1, 2)
         return x
 
 

lucid/models/imgclf/maxvit.py

@@ -216,7 +216,7 @@ def _grid_reverse(
 
 def _get_relative_position_index(win_h: int, win_w: int) -> Tensor:
     coords = lucid.stack(lucid.meshgrid(lucid.arange(win_h), lucid.arange(win_w)))
-    coords_flatten = lucid.flatten(coords, axis=1)
+    coords_flatten = lucid.flatten(coords, start_axis=1)
 
     relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
     relative_coords = relative_coords.transpose((1, 2, 0))

lucid/models/imgclf/pvt.py

@@ -328,7 +328,7 @@ class _DWConv(nn.Module):
         B, _, C = x.shape
         x = x.swapaxes(1, 2).reshape(B, C, H, W)
         x = self.dwconv(x)
-        x = x.flatten(axis=2).swapaxes(1, 2)
+        x = x.flatten(start_axis=2).swapaxes(1, 2)
 
         return x
 
@@ -548,7 +548,7 @@ class _OverlapPatchEmbed(nn.Module):
     def forward(self, x: Tensor) -> tuple[Tensor, int, int]:
         x = self.proj(x)
         H, W = x.shape[2:]
-        x = x.flatten(axis=2).swapaxes(1, 2)
+        x = x.flatten(start_axis=2).swapaxes(1, 2)
         x = self.norm(x)
 
         return x, H, W

lucid/models/imggen/vae.py

@@ -51,7 +51,7 @@ class VAE(nn.Module):
         h = x
         for encoder in self.encoders:
             h = encoder(h)
-            mu, logvar = lucid.split(h, 2, axis=1)
+            mu, logvar = lucid.chunk(h, 2, axis=1)
             z = self.reparameterize(mu, logvar)
 
             mus.append(mu)

lucid/models/objdet/efficientdet.py

@@ -82,23 +82,37 @@ class _BiFPN(nn.Module):
     def _norm_weight(self, weight: Tensor) -> Tensor:
         return weight / (weight.sum(axis=0) + self.eps)
 
+    @staticmethod
+    def _resize_like(x: Tensor, ref: Tensor) -> Tensor:
+        if x.shape[2:] == ref.shape[2:]:
+            return x
+        return F.interpolate(x, size=ref.shape[2:], mode="nearest")
+
     def _forward_up(self, feats: tuple[Tensor]) -> tuple[Tensor]:
         p3_in, p4_in, p5_in, p6_in, p7_in = feats
 
         w1_p6_up = self._norm_weight(self.acts["6_w1"](self.weights["6_w1"]))
-        p6_up_in = w1_p6_up[0] * p6_in + w1_p6_up[1] * self.ups["6"](p7_in)
+        p6_up_in = w1_p6_up[0] * p6_in + w1_p6_up[1] * self._resize_like(
+            self.ups["6"](p7_in), p6_in
+        )
         p6_up = self.convs["6_up"](p6_up_in)
 
         w1_p5_up = self._norm_weight(self.acts["5_w1"](self.weights["5_w1"]))
-        p5_up_in = w1_p5_up[0] * p5_in + w1_p5_up[1] * self.ups["5"](p6_up)
+        p5_up_in = w1_p5_up[0] * p5_in + w1_p5_up[1] * self._resize_like(
+            self.ups["5"](p6_up), p5_in
+        )
         p5_up = self.convs["5_up"](p5_up_in)
 
         w1_p4_up = self._norm_weight(self.acts["4_w1"](self.weights["4_w1"]))
-        p4_up_in = w1_p4_up[0] * p4_in + w1_p4_up[1] * self.ups["4"](p5_up)
+        p4_up_in = w1_p4_up[0] * p4_in + w1_p4_up[1] * self._resize_like(
+            self.ups["4"](p5_up), p4_in
+        )
         p4_up = self.convs["4_up"](p4_up_in)
 
         w1_p3_up = self._norm_weight(self.acts["3_w1"](self.weights["3_w1"]))
-        p3_up_in = w1_p3_up[0] * p3_in + w1_p3_up[1] * self.ups["3"](p4_up)
+        p3_up_in = w1_p3_up[0] * p3_in + w1_p3_up[1] * self._resize_like(
+            self.ups["3"](p4_up), p3_in
+        )
         p3_out = self.convs["3_up"](p3_up_in)
 
         return p3_out, p4_up, p5_up, p6_up
@@ -113,7 +127,7 @@ class _BiFPN(nn.Module):
         p4_down_in = (
             w2_p4_down[0] * p4_in
             + w2_p4_down[1] * p4_up
-            + w2_p4_down[2] * self.downs["4"](p3_out)
+            + w2_p4_down[2] * self._resize_like(self.downs["4"](p3_out), p4_in)
         )
         p4_out = self.convs["4_down"](p4_down_in)
 
@@ -121,7 +135,7 @@ class _BiFPN(nn.Module):
         p5_down_in = (
             w2_p5_down[0] * p5_in
             + w2_p5_down[1] * p5_up
-            + w2_p5_down[2] * self.downs["5"](p4_out)
+            + w2_p5_down[2] * self._resize_like(self.downs["5"](p4_out), p5_in)
         )
         p5_out = self.convs["5_down"](p5_down_in)
 
@@ -129,12 +143,14 @@ class _BiFPN(nn.Module):
         p6_down_in = (
             w2_p6_down[0] * p6_in
             + w2_p6_down[1] * p6_up
-            + w2_p6_down[2] * self.downs["6"](p5_out)
+            + w2_p6_down[2] * self._resize_like(self.downs["6"](p5_out), p6_in)
         )
         p6_out = self.convs["6_down"](p6_down_in)
 
         w2_p7_down = self._norm_weight(self.acts["7_w2"](self.weights["7_w2"]))
-        p7_down_in = w2_p7_down[0] * p7_in + w2_p7_down[1] * self.downs["7"](p6_out)
+        p7_down_in = w2_p7_down[0] * p7_in + w2_p7_down[1] * self._resize_like(
+            self.downs["7"](p6_out), p7_in
+        )
         p7_out = self.convs["7_down"](p7_down_in)
 
         return p3_out, p4_out, p5_out, p6_out, p7_out

lucid/models/objdet/rcnn.py

@@ -120,7 +120,7 @@ class RCNN(nn.Module):
 
         if isinstance(feats, (tuple, list)):
             feats = feats[-1]
-        feats = feats.flatten(axis=1)
+        feats = feats.flatten(start_axis=1)
 
         cls_scores = self.svm(feats)
         bbox_deltas = self.bbox_reg(feats)

lucid/models/objdet/util.py

@@ -283,6 +283,11 @@ class SelectiveSearch(nn.Module):
 
 
 def iou(boxes_a: Tensor, boxes_b: Tensor) -> Tensor:
+    if boxes_a.ndim == 1:
+        boxes_a = boxes_a.unsqueeze(0)
+    if boxes_b.ndim == 1:
+        boxes_b = boxes_b.unsqueeze(0)
+
     x1a, y1a, x2a, y2a = boxes_a.unbind(axis=1)
     x1b, y1b, x2b, y2b = boxes_b.unbind(axis=1)
 

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
 
@@ -29,9 +44,6 @@ __all__ = [
     "set_state_dict_pass_attr",
 ]
 
-_ForwardHookType = Callable[["Module", tuple[Tensor], tuple[Tensor]], None]
-_BackwardHookType = Callable[[Tensor, _NumPyArray], None]
-
 
 class Module:
     _registry_map: dict[Type, OrderedDict[str, Any]] = {}
@@ -49,8 +61,20 @@ class Module:
         self.training = True
         self.device: _DeviceType = "cpu"
 
-        self._forward_hooks: list[_ForwardHookType] = []
-        self._backward_hooks: list[_BackwardHookType] = []
+        self._forward_pre_hooks: list[
+            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()
 
@@ -106,14 +130,53 @@ class Module:
 
         self.__setattr__(name, buffer)
 
-    def register_forward_hook(self, hook: _ForwardHookType) -> Callable:
-        self._forward_hooks.append(hook)
-        return lambda: self._forward_hooks.remove(hook)
-
-    def register_backward_hook(self, hook: _BackwardHookType) -> Callable:
+    def register_forward_pre_hook(
+        self,
+        hook: _ForwardPreHook | _ForwardPreHookKwargs,
+        *,
+        with_kwargs: bool = False,
+    ) -> Callable:
+        self._forward_pre_hooks.append((hook, with_kwargs))
+        return lambda: self._forward_pre_hooks.remove((hook, with_kwargs))
+
+    def register_forward_hook(
+        self, hook: _ForwardHook | _ForwardHookKwargs, *, with_kwargs: bool = False
+    ) -> Callable:
+        self._forward_hooks.append((hook, with_kwargs))
+        return lambda: self._forward_hooks.remove((hook, with_kwargs))
+
+    def register_backward_hook(self, hook: _BackwardHook) -> Callable:
         self._backward_hooks.append(hook)
         return lambda: self._backward_hooks.remove(hook)
 
+    def register_full_backward_pre_hook(self, hook: _FullBackwardPreHook) -> Callable:
+        self._full_backward_pre_hooks.append(hook)
+        return lambda: self._full_backward_pre_hooks.remove(hook)
+
+    def register_full_backward_hook(self, hook: _FullBackwardHook) -> Callable:
+        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()
@@ -190,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()
 
@@ -208,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())
@@ -236,15 +308,72 @@ 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:
+                result = hook(self, args, kwargs)
+                if result is not None:
+                    args, kwargs = result
+            else:
+                result = hook(self, args)
+                if result is not None:
+                    args = result
+
         output = self.forward(*args, **kwargs)
-        for hook in self._forward_hooks:
-            hook(self, args, output)
+
+        for hook, with_kwargs in self._forward_hooks:
+            if with_kwargs:
+                result = hook(self, args, kwargs, output)
+            else:
+                result = hook(self, args, output)
+            if result is not None:
+                output = result
 
         if isinstance(output, Tensor) and self._backward_hooks:
             for hook in self._backward_hooks:
                 output.register_hook(hook)
 
+        if self._full_backward_pre_hooks or self._full_backward_hooks:
+            outputs = output if isinstance(output, tuple) else (output,)
+            output_tensors = [out for out in outputs if isinstance(out, Tensor)]
+
+            if output_tensors:
+                grad_outputs: list[_NumPyArray | None] = [None] * len(output_tensors)
+                called = False
+
+                def _call_full_backward_hooks() -> None:
+                    nonlocal called, grad_outputs
+                    if called:
+                        return
+                    called = True
+
+                    grad_output_tuple = tuple(grad_outputs)
+                    for hook in self._full_backward_pre_hooks:
+                        result = hook(self, grad_output_tuple)
+                        if result is not None:
+                            grad_output_tuple = result
+
+                    grad_input_tuple = tuple(
+                        arg.grad if isinstance(arg, Tensor) else None for arg in args
+                    )
+                    for hook in self._full_backward_hooks:
+                        hook(self, grad_input_tuple, grad_output_tuple)
+
+                for idx, out in enumerate(output_tensors):
+
+                    def _make_hook(index: int) -> Callable:
+                        def _hook(_, grad: _NumPyArray) -> None:
+                            grad_outputs[index] = grad
+                            if all(g is not None for g in grad_outputs):
+                                _call_full_backward_hooks()
+
+                        return _hook
+
+                    out.register_hook(_make_hook(idx))
+
         return output
 
     def __repr__(self) -> str:

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/visual/__init__.py

@@ -1 +1,2 @@
 from .graph import *
+from .mermaid import *

lucid/visual/graph.py

@@ -1,4 +1,6 @@
 from typing import Union
+from warnings import deprecated
+
 import networkx as nx
 import matplotlib.pyplot as plt
 
@@ -9,6 +11,7 @@ from lucid._tensor import Tensor
 __all__ = ["draw_tensor_graph"]
 
 
+@deprecated("This feature will be re-written with Mermaid in future relases.")
 def draw_tensor_graph(
     tensor: Tensor,
     horizontal: bool = False,