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

lucid/datasets/__init__.py

@@ -1,3 +1,2 @@
-# TODO: implement per-batch data loading and transformations
 from .mnist import *
 from .cifar import *

lucid/datasets/cifar.py

@@ -4,6 +4,8 @@ import openml
 import math
 
 from typing import SupportsIndex, Tuple, ClassVar
+from pathlib import Path
+import re
 
 import lucid
 from lucid._tensor import Tensor
@@ -17,6 +19,42 @@ __all__ = ["CIFAR10", "CIFAR100"]
 class CIFAR10(DatasetBase):
     OPENML_ID: ClassVar[int] = 40927
 
+    def __init__(
+        self,
+        root: str | Path,
+        train: bool | None = True,
+        download: bool | None = False,
+        transform: lucid.nn.Module | None = None,
+        target_transform: lucid.nn.Module | None = None,
+        test_size: float = 0.2,
+        to_tensor: bool = True,
+        *,
+        cache: bool = True,
+        scale: float | None = None,
+        resize: tuple[int, int] | None = None,
+        normalize: tuple[tuple[float, ...], tuple[float, ...]] | None = None,
+        cache_preprocessed: bool = True,
+        preprocess_dtype: lucid.Numeric = lucid.Float16,
+        preprocess_chunk_size: int = 4096,
+    ) -> None:
+        self.cache = cache
+        self.scale = scale
+        self.resize = resize
+        self.normalize = normalize
+        self.cache_preprocessed = cache_preprocessed
+        self.preprocess_dtype = preprocess_dtype
+        self.preprocess_chunk_size = preprocess_chunk_size
+
+        super().__init__(
+            root=root,
+            train=train,
+            download=download,
+            transform=transform,
+            target_transform=target_transform,
+            test_size=test_size,
+            to_tensor=to_tensor,
+        )
+
     def _download(self) -> None:
         try:
             dataset = openml.datasets.get_dataset(self.OPENML_ID)
@@ -26,7 +64,36 @@ class CIFAR10(DatasetBase):
         except Exception as e:
             raise RuntimeError(f"Failed to download the CIFAR-10 dataset. Error: {e}")
 
-    def _load_data(self, split: str) -> Tuple[Tensor, Tensor]:
+    def _cache_key(self) -> str:
+        parts: list[str] = []
+        if self.scale is not None:
+            parts.append(f"s{self.scale:g}")
+        if self.resize is not None:
+            parts.append(f"r{self.resize[0]}x{self.resize[1]}")
+        if self.normalize is not None:
+            mean, std = self.normalize
+            parts.append("m" + ",".join(f"{v:g}" for v in mean))
+            parts.append("v" + ",".join(f"{v:g}" for v in std))
+        if not parts:
+            return "raw"
+        key = "_".join(parts)
+        return re.sub(r"[^a-zA-Z0-9_,.x-]+", "_", key)
+
+    def _raw_cache_path(self) -> Path:
+        return self.root / "CIFAR10_uint8.npz"
+
+    def _proc_cache_path(self) -> Path:
+        dtype_name = str(self.preprocess_dtype)
+        return self.root / f"CIFAR10_{self._cache_key()}_{dtype_name}.npz"
+
+    def _ensure_raw_cache(self) -> tuple[np.ndarray, np.ndarray]:
+        raw_path = self._raw_cache_path()
+        if self.cache and raw_path.exists():
+            with np.load(raw_path) as npz:
+                images = npz["images"]
+                labels = npz["labels"]
+            return images, labels
+
         csv_path = self.root / "CIFAR10.csv"
         if not csv_path.exists():
             raise RuntimeError(
@@ -36,9 +103,69 @@ class CIFAR10(DatasetBase):
 
         df = pd.read_csv(csv_path)
         labels = df["class"].values.astype(np.int32)
-        images = df.drop(columns=["class"]).values.astype(np.float32)
+        images = df.drop(columns=["class"]).values.astype(np.uint8, copy=False)
         images = images.reshape(-1, 3, 32, 32)
 
+        if self.cache:
+            np.savez_compressed(raw_path, images=images, labels=labels)
+
+        return images, labels
+
+    def _maybe_preprocess_and_cache(
+        self, images_uint8: np.ndarray, labels_int32: np.ndarray
+    ) -> tuple[np.ndarray, np.ndarray]:
+        if self.resize is None and self.scale is None and self.normalize is None:
+            return images_uint8.astype(np.float32), labels_int32
+
+        proc_path = self._proc_cache_path()
+        if self.cache and self.cache_preprocessed and proc_path.exists():
+            with np.load(proc_path) as npz:
+                images = npz["images"]
+                labels = npz["labels"]
+            return images, labels
+
+        from lucid.transforms import Compose, Resize, Normalize
+
+        class _Scale(lucid.nn.Module):
+            def __init__(self, factor: float) -> None:
+                super().__init__()
+                self.factor = factor
+
+            def forward(self, x: Tensor) -> Tensor:
+                return x * self.factor
+
+        transforms: list[lucid.nn.Module] = []
+        if self.resize is not None:
+            transforms.append(Resize(self.resize))
+        if self.scale is not None:
+            transforms.append(_Scale(self.scale))
+        if self.normalize is not None:
+            mean, std = self.normalize
+            transforms.append(Normalize(mean=mean, std=std))
+
+        transform = Compose(transforms)
+        n = images_uint8.shape[0]
+        out_h, out_w = self.resize if self.resize is not None else (32, 32)
+
+        out_dtype = np.float16 if self.preprocess_dtype == lucid.Float16 else np.float32
+        out_images = np.empty((n, 3, out_h, out_w), dtype=out_dtype)
+
+        for start in range(0, n, self.preprocess_chunk_size):
+            end = min(start + self.preprocess_chunk_size, n)
+            chunk = images_uint8[start:end].astype(np.float32)
+            x = lucid.to_tensor(chunk, dtype=lucid.Float32)
+            x = transform(x)
+            out_images[start:end] = x.numpy().astype(out_dtype, copy=False)
+
+        if self.cache and self.cache_preprocessed:
+            np.savez_compressed(proc_path, images=out_images, labels=labels_int32)
+
+        return out_images, labels_int32
+
+    def _load_data(self, split: str) -> Tuple[Tensor, Tensor]:
+        images, labels = self._ensure_raw_cache()
+        images, labels = self._maybe_preprocess_and_cache(images, labels)
+
         train_size = int(math.floor(len(images) * (1 - self.test_size)))
         if split == "train":
             images, labels = images[:train_size], labels[:train_size]
@@ -52,7 +179,7 @@ class CIFAR10(DatasetBase):
         return images, labels
 
     def __getitem__(self, index: SupportsIndex) -> Tuple[Tensor, Tensor]:
-        image = self.data[index].reshape(-1, 3, 32, 32)
+        image = self.data[index]
         label = self.targets[index]
 
         if self.transform:
@@ -66,6 +193,42 @@ class CIFAR10(DatasetBase):
 class CIFAR100(DatasetBase):
     OPENML_ID: ClassVar[int] = 41983
 
+    def __init__(
+        self,
+        root: str | Path,
+        train: bool | None = True,
+        download: bool | None = False,
+        transform: lucid.nn.Module | None = None,
+        target_transform: lucid.nn.Module | None = None,
+        test_size: float = 0.2,
+        to_tensor: bool = True,
+        *,
+        cache: bool = True,
+        scale: float | None = None,
+        resize: tuple[int, int] | None = None,
+        normalize: tuple[tuple[float, ...], tuple[float, ...]] | None = None,
+        cache_preprocessed: bool = True,
+        preprocess_dtype: lucid.Numeric = lucid.Float16,
+        preprocess_chunk_size: int = 4096,
+    ) -> None:
+        self.cache = cache
+        self.scale = scale
+        self.resize = resize
+        self.normalize = normalize
+        self.cache_preprocessed = cache_preprocessed
+        self.preprocess_dtype = preprocess_dtype
+        self.preprocess_chunk_size = preprocess_chunk_size
+
+        super().__init__(
+            root=root,
+            train=train,
+            download=download,
+            transform=transform,
+            target_transform=target_transform,
+            test_size=test_size,
+            to_tensor=to_tensor,
+        )
+
     def _download(self) -> None:
         try:
             dataset = openml.datasets.get_dataset(self.OPENML_ID)
@@ -75,7 +238,37 @@ class CIFAR100(DatasetBase):
         except Exception as e:
             raise RuntimeError(f"Failed to download the CIFAR-100 dataset. Error: {e}")
 
-    def _load_data(self, split: str) -> Tuple[Tensor, Tensor]:
+    def _cache_key(self) -> str:
+        parts: list[str] = []
+        if self.scale is not None:
+            parts.append(f"s{self.scale:g}")
+        if self.resize is not None:
+            parts.append(f"r{self.resize[0]}x{self.resize[1]}")
+        if self.normalize is not None:
+            mean, std = self.normalize
+            parts.append("m" + ",".join(f"{v:g}" for v in mean))
+            parts.append("v" + ",".join(f"{v:g}" for v in std))
+        if not parts:
+            return "raw"
+
+        key = "_".join(parts)
+        return re.sub(r"[^a-zA-Z0-9_,.x-]+", "_", key)
+
+    def _raw_cache_path(self) -> Path:
+        return self.root / "CIFAR100_uint8.npz"
+
+    def _proc_cache_path(self) -> Path:
+        dtype_name = str(self.preprocess_dtype)
+        return self.root / f"CIFAR100_{self._cache_key()}_{dtype_name}.npz"
+
+    def _ensure_raw_cache(self) -> tuple[np.ndarray, np.ndarray]:
+        raw_path = self._raw_cache_path()
+        if self.cache and raw_path.exists():
+            with np.load(raw_path) as npz:
+                images = npz["images"]
+                labels = npz["labels"]
+            return images, labels
+
         csv_path = self.root / "CIFAR100.csv"
         if not csv_path.exists():
             raise RuntimeError(
@@ -85,9 +278,69 @@ class CIFAR100(DatasetBase):
 
         df = pd.read_csv(csv_path)
         labels = df["class"].values.astype(np.int32)
-        images = df.drop(columns=["class"]).values.astype(np.float32)
+        images = df.drop(columns=["class"]).values.astype(np.uint8, copy=False)
         images = images.reshape(-1, 3, 32, 32)
 
+        if self.cache:
+            np.savez_compressed(raw_path, images=images, labels=labels)
+
+        return images, labels
+
+    def _maybe_preprocess_and_cache(
+        self, images_uint8: np.ndarray, labels_int32: np.ndarray
+    ) -> tuple[np.ndarray, np.ndarray]:
+        if self.resize is None and self.scale is None and self.normalize is None:
+            return images_uint8.astype(np.float32), labels_int32
+
+        proc_path = self._proc_cache_path()
+        if self.cache and self.cache_preprocessed and proc_path.exists():
+            with np.load(proc_path) as npz:
+                images = npz["images"]
+                labels = npz["labels"]
+            return images, labels
+
+        from lucid.transforms import Compose, Resize, Normalize
+
+        class _Scale(lucid.nn.Module):
+            def __init__(self, factor: float) -> None:
+                super().__init__()
+                self.factor = factor
+
+            def forward(self, x: Tensor) -> Tensor:
+                return x * self.factor
+
+        transforms: list[lucid.nn.Module] = []
+        if self.resize is not None:
+            transforms.append(Resize(self.resize))
+        if self.scale is not None:
+            transforms.append(_Scale(self.scale))
+        if self.normalize is not None:
+            mean, std = self.normalize
+            transforms.append(Normalize(mean=mean, std=std))
+
+        transform = Compose(transforms)
+        n = images_uint8.shape[0]
+        out_h, out_w = self.resize if self.resize is not None else (32, 32)
+
+        out_dtype = np.float16 if self.preprocess_dtype == lucid.Float16 else np.float32
+        out_images = np.empty((n, 3, out_h, out_w), dtype=out_dtype)
+
+        for start in range(0, n, self.preprocess_chunk_size):
+            end = min(start + self.preprocess_chunk_size, n)
+            chunk = images_uint8[start:end].astype(np.float32)
+            x = lucid.to_tensor(chunk, dtype=lucid.Float32)
+            x = transform(x)
+            out_images[start:end] = x.numpy().astype(out_dtype, copy=False)
+
+        if self.cache and self.cache_preprocessed:
+            np.savez_compressed(proc_path, images=out_images, labels=labels_int32)
+
+        return out_images, labels_int32
+
+    def _load_data(self, split: str) -> Tuple[Tensor, Tensor]:
+        images, labels = self._ensure_raw_cache()
+        images, labels = self._maybe_preprocess_and_cache(images, labels)
+
         train_size = int(math.floor(len(images) * (1 - self.test_size)))
         if split == "train":
             images, labels = images[:train_size], labels[:train_size]
@@ -101,7 +354,7 @@ class CIFAR100(DatasetBase):
         return images, labels
 
     def __getitem__(self, index: SupportsIndex) -> Tuple[Tensor, Tensor]:
-        image = self.data[index].reshape(-1, 3, 32, 32)
+        image = self.data[index]
         label = self.targets[index]
 
         if self.transform:

lucid/models/imgclf/vit.py

@@ -32,10 +32,12 @@ class ViT(nn.Module):
             in_channels, embedding_dim, kernel_size=patch_size, stride=patch_size
         )
 
-        self.cls_token = nn.Parameter(lucid.random.randn(1, 1, embedding_dim))
-        self.pos_emb = nn.Parameter(
-            lucid.random.randn(1, 1 + self.num_patches, embedding_dim)
-        )
+        self.cls_token = nn.Parameter(lucid.zeros(1, 1, embedding_dim))
+        self.pos_emb = nn.Parameter(lucid.zeros(1, 1 + self.num_patches, embedding_dim))
+
+        nn.init.normal(self.cls_token, std=0.02)
+        nn.init.normal(self.pos_emb, std=0.02)
+
         self.dropout = nn.Dropout(dropout_rate)
 
         encoder_layer = nn.TransformerEncoderLayer(

lucid/nn/__init__.py

@@ -4,4 +4,4 @@ from lucid.nn.modules import *
 from lucid.nn.fused import *
 
 import lucid.nn.init as init
-import lucid.nn.util as util
+import lucid.nn.utils as utils

lucid/nn/modules/rnn.py

@@ -5,6 +5,7 @@ import lucid.nn as nn
 import lucid.nn.functional as F
 
 from lucid._tensor import Tensor
+from lucid.nn.utils.rnn import PackedSequence
 from lucid.types import Numeric, _DeviceType
 
 from .activation import Tanh, ReLU
@@ -351,21 +352,47 @@ class RNNBase(nn.Module):
         )
 
     def forward(
-        self, input_: Tensor, hx: Tensor | tuple[Tensor, Tensor] | None = None
-    ) -> tuple[Tensor, Tensor] | tuple[Tensor, tuple[Tensor, Tensor]]:
-        if input_.ndim != 3:
-            raise ValueError(
-                f"RNNBase expected input with 3 dimensions, got {input_.ndim} dimensions"
-            )
+        self,
+        input_: Tensor | PackedSequence,
+        hx: Tensor | tuple[Tensor, Tensor] | None = None,
+    ) -> (
+        tuple[Tensor | PackedSequence, Tensor]
+        | tuple[Tensor | PackedSequence, tuple[Tensor, Tensor]]
+    ):
+        is_packed = isinstance(input_, PackedSequence)
+        if is_packed:
+            data = input_.data
+            batch_sizes = input_.batch_sizes
+            if data.ndim != 2:
+                raise ValueError(
+                    "RNNBase expected packed data with 2 dimensions, "
+                    f"got {data.ndim} dimensions"
+                )
+            if batch_sizes.ndim != 1 or batch_sizes.shape[0] == 0:
+                raise ValueError(
+                    "PackedSequence batch_sizes must be a non-empty 1D tensor"
+                )
 
-        if self.batch_first:
-            input_ = input_.swapaxes(0, 1)
+            batch_size = int(batch_sizes[0].item())
+            feat = data.shape[1]
+            if feat != self.input_size:
+                raise ValueError(
+                    f"RNNBase expected input with feature size {self.input_size}, got {feat}"
+                )
+        else:
+            if input_.ndim != 3:
+                raise ValueError(
+                    f"RNNBase expected input with 3 dimensions, got {input_.ndim} dimensions"
+                )
 
-        seq_len, batch_size, feat = input_.shape
-        if feat != self.input_size:
-            raise ValueError(
-                f"RNNBase expected input with feature size {self.input_size}, got {feat}"
-            )
+            if self.batch_first:
+                input_ = input_.swapaxes(0, 1)
+
+            seq_len, batch_size, feat = input_.shape
+            if feat != self.input_size:
+                raise ValueError(
+                    f"RNNBase expected input with feature size {self.input_size}, got {feat}"
+                )
 
         if self.is_lstm:
             if hx is None:
@@ -410,7 +437,7 @@ class RNNBase(nn.Module):
             if hx.shape[2] != self.hidden_size:
                 raise ValueError("Incorrect hidden size in hx")
 
-        layer_input = input_
+        layer_input = data if is_packed else input_
         h_n_list: list[Tensor] = []
         c_n_list: list[Tensor] | None = [] if self.is_lstm else None
 
@@ -420,33 +447,111 @@ class RNNBase(nn.Module):
                 c_t = hx_c[layer_idx]
             else:
                 h_t = hx[layer_idx]
+
             outputs = []
+            if is_packed:
+                final_h: list[Tensor] = []
+                final_c: list[Tensor] | None = [] if self.is_lstm else None
+                offset = 0
+
+                prev_bs: int | None = None
+                max_len = int(batch_sizes.shape[0])
+                for t in range(max_len):
+                    bs = int(batch_sizes[t].item())
+                    if bs == 0:
+                        break
+
+                    if prev_bs is None:
+                        prev_bs = bs
+                    if bs > prev_bs:
+                        raise ValueError(
+                            "PackedSequence batch_sizes must be non-increasing"
+                        )
+
+                    if bs < prev_bs:
+                        final_h.append(h_t[bs:prev_bs])
+                        if self.is_lstm and final_c is not None:
+                            final_c.append(c_t[bs:prev_bs])
+
+                        h_t = h_t[:bs]
+                        if self.is_lstm:
+                            c_t = c_t[:bs]
+
+                    step_input = layer_input[offset : offset + bs]
+                    offset += bs
+
+                    if self.is_lstm:
+                        h_t, c_t = cell(step_input, (h_t, c_t))
+                    else:
+                        h_t = cell(step_input, h_t)
+
+                    outputs.append(h_t)
+                    prev_bs = bs
+
+                final_h.append(h_t)
+                if self.is_lstm and final_c is not None:
+                    final_c.append(c_t)
+
+                h_n_list.append(
+                    lucid.concatenate(tuple(reversed(final_h)), axis=0).unsqueeze(
+                        axis=0
+                    )
+                )
+                if self.is_lstm and final_c is not None and c_n_list is not None:
+                    c_n_list.append(
+                        lucid.concatenate(tuple(reversed(final_c)), axis=0).unsqueeze(
+                            axis=0
+                        )
+                    )
+
+                layer_output = (
+                    lucid.concatenate(tuple(outputs), axis=0)
+                    if outputs
+                    else layer_input[:0]
+                )
 
-            for t in range(seq_len):
-                if self.is_lstm:
-                    h_t, c_t = cell(layer_input[t], (h_t, c_t))
-                    outputs.append(h_t.unsqueeze(axis=0))
-                else:
-                    h_t = cell(layer_input[t], h_t)
-                    outputs.append(h_t.unsqueeze(axis=0))
+            else:
+                for t in range(seq_len):
+                    if self.is_lstm:
+                        h_t, c_t = cell(layer_input[t], (h_t, c_t))
+                        outputs.append(h_t.unsqueeze(axis=0))
+                    else:
+                        h_t = cell(layer_input[t], h_t)
+                        outputs.append(h_t.unsqueeze(axis=0))
 
-            layer_output = lucid.concatenate(tuple(outputs), axis=0)
+                layer_output = lucid.concatenate(tuple(outputs), axis=0)
 
             if self.training and self.dropout > 0.0 and layer_idx < self.num_layers - 1:
                 layer_output = F.dropout(layer_output, p=self.dropout)
 
-            h_n_list.append(h_t.unsqueeze(axis=0))
-            if self.is_lstm and c_n_list is not None:
-                c_n_list.append(c_t.unsqueeze(axis=0))
+            if not is_packed:
+                h_n_list.append(h_t.unsqueeze(axis=0))
+                if self.is_lstm and c_n_list is not None:
+                    c_n_list.append(c_t.unsqueeze(axis=0))
             layer_input = layer_output
 
-        output = layer_input
+        if is_packed:
+            output = PackedSequence(
+                data=layer_input,
+                batch_sizes=batch_sizes,
+                sorted_indices=input_.sorted_indices,
+                unsorted_indices=input_.unsorted_indices,
+            )
+        else:
+            output = layer_input
+
         h_n = lucid.concatenate(tuple(h_n_list), axis=0)
         if self.is_lstm and c_n_list is not None:
             c_n = lucid.concatenate(tuple(c_n_list), axis=0)
 
-        if self.batch_first:
-            output = output.swapaxes(0, 1)
+        if is_packed:
+            if input_.unsorted_indices is not None:
+                h_n = h_n[:, input_.unsorted_indices]
+                if self.is_lstm and c_n_list is not None:
+                    c_n = c_n[:, input_.unsorted_indices]
+        else:
+            if self.batch_first:
+                output = output.swapaxes(0, 1)
 
         if self.is_lstm and c_n_list is not None:
             return output, (h_n, c_n)

lucid/nn/utils/__init__.py

@@ -0,0 +1,2 @@
+from ._grad import *
+from . import rnn as rnn

lucid/nn/{util.py → utils/_grad.py}

@@ -6,7 +6,7 @@ from lucid._tensor import Tensor
 from lucid.types import _Scalar
 
 
-__all__ = ["grad_norm", "clip_grad_norm", "clip_grad_value"]
+__all__ = ["grad_norm", "get_total_norm", "clip_grad_norm", "clip_grad_value"]
 
 
 def _as_iter(parameters: Iterable[Tensor] | Tensor) -> list[Tensor]:
@@ -32,6 +32,25 @@ def grad_norm(parameters: Iterable[Tensor] | Tensor, norm_type: int = 2) -> Tens
     return Tensor(total_norm, device=device)
 
 
+def get_total_norm(parameters: Iterable[Tensor] | Tensor, norm_type: int = 2) -> Tensor:
+    parameters = _as_iter(parameters)
+    if not parameters:
+        return Tensor(0.0)
+
+    device = parameters[0].device
+    grads: list[Tensor] = [p.grad for p in parameters if p.grad is not None]
+    if not grads:
+        return Tensor(0.0, device=device)
+
+    norm_pow_sum = 0.0
+    for g in grads:
+        grad_norm = lucid.linalg.norm(lucid.ravel(g), ord=norm_type).item()
+        norm_pow_sum += grad_norm**norm_type
+
+    total_norm = norm_pow_sum ** (1.0 / norm_type)
+    return Tensor(total_norm, device=device)
+
+
 def clip_grad_norm(
     parameters: Iterable[Tensor] | Tensor,
     max_norm: _Scalar,
@@ -39,7 +58,7 @@ def clip_grad_norm(
     eps: float = 1e-7,
 ) -> float:
     params: list[Tensor] = [p for p in _as_iter(parameters) if p.grad is not None]
-    total_norm = grad_norm(params, norm_type=norm_type)
+    total_norm = get_total_norm(params, norm_type=norm_type)
 
     clip_coef = float(max_norm) / (total_norm.item() + eps)
     if clip_coef < 1.0:

lucid/nn/utils/rnn.py

@@ -0,0 +1,237 @@
+from dataclasses import dataclass
+from typing import Iterable, Sequence
+
+import lucid
+
+from lucid._tensor import Tensor
+from lucid.types import _Scalar
+
+
+__all__ = [
+    "PackedSequence",
+    "pad_sequence",
+    "pack_padded_sequence",
+    "pad_packed_sequence",
+    "pack_sequence",
+    "unpack_sequence",
+]
+
+
+@dataclass(frozen=True)
+class PackedSequence:
+    data: Tensor
+    batch_sizes: Tensor
+    sorted_indices: Tensor | None = None
+    unsorted_indices: Tensor | None = None
+
+
+def pad_sequence(
+    sequences: Iterable[Tensor], batch_first: bool = False, padding_value: _Scalar = 0
+) -> Tensor:
+    seq_list = list(sequences)
+    if not seq_list:
+        raise ValueError("pad_sequence expected a non-empty iterable of Tensors")
+
+    first = seq_list[0]
+    if not isinstance(first, Tensor):
+        raise TypeError("pad_sequence expects Tensor elements")
+
+    ndim = first.ndim
+    if ndim < 1:
+        raise ValueError("pad_sequence expects tensors with at least 1 dimension")
+
+    trailing_shape = first.shape[1:]
+    device = first.device
+    dtype = first.dtype
+
+    lengths: list[int] = []
+    for idx, seq in enumerate(seq_list):
+        if not isinstance(seq, Tensor):
+            raise TypeError("pad_sequence expects Tensor elements")
+        if seq.ndim != ndim:
+            raise ValueError(
+                f"pad_sequence expects tensors with {ndim} dimensions, "
+                f"got {seq.ndim} at index {idx}"
+            )
+        if seq.shape[1:] != trailing_shape:
+            raise ValueError(
+                "pad_sequence expects all tensors to share the same trailing shape"
+            )
+        if seq.device != device:
+            raise ValueError("pad_sequence expects all tensors on the same device")
+        if seq.dtype != dtype:
+            raise ValueError("pad_sequence expects all tensors with the same dtype")
+        lengths.append(seq.shape[0])
+
+    max_len = max(lengths)
+    batch_size = len(seq_list)
+
+    if batch_first:
+        out_shape = (batch_size, max_len, *trailing_shape)
+    else:
+        out_shape = (max_len, batch_size, *trailing_shape)
+
+    output = lucid.full(out_shape, padding_value, dtype=dtype, device=device)
+    for i, seq in enumerate(seq_list):
+        length = lengths[i]
+        if length == 0:
+            continue
+        if batch_first:
+            output[i, :length] = seq
+        else:
+            output[:length, i] = seq
+
+    return output
+
+
+def _as_lengths(lengths: Sequence[int] | Tensor, *, device: str) -> Tensor:
+    if isinstance(lengths, Tensor):
+        return lengths
+    return Tensor(list(lengths), device=device)
+
+
+def _invert_permutation(indices: Tensor) -> Tensor:
+    return lucid.argsort(indices, axis=0)
+
+
+def pack_padded_sequence(
+    input_: Tensor,
+    lengths: Sequence[int] | Tensor,
+    batch_first: bool = False,
+    enforce_sorted: bool = True,
+) -> PackedSequence:
+    if input_.ndim < 2:
+        raise ValueError(
+            f"pack_padded_sequence expected input with at least 2 dims, got {input_.ndim}"
+        )
+
+    if batch_first:
+        input_ = input_.swapaxes(0, 1)
+
+    seq_len, batch_size = input_.shape[0], input_.shape[1]
+    lengths_t = _as_lengths(lengths, device=input_.device)
+    if lengths_t.ndim != 1:
+        raise ValueError("lengths must be a 1D sequence or tensor")
+    if lengths_t.shape[0] != batch_size:
+        raise ValueError(
+            f"lengths size {lengths_t.shape[0]} does not match batch size {batch_size}"
+        )
+
+    sorted_indices = None
+    unsorted_indices = None
+
+    if enforce_sorted:
+        sorted_lengths = lengths_t
+    else:
+        sorted_indices = lucid.argsort(lengths_t, descending=True, axis=0)
+        unsorted_indices = _invert_permutation(sorted_indices)
+
+        lengths_t = lengths_t[sorted_indices]
+        input_ = input_[:, sorted_indices]
+        sorted_lengths = lengths_t
+
+    max_len = int(sorted_lengths[0].item())
+    if max_len > seq_len:
+        raise ValueError(
+            f"lengths has max {max_len} but input has sequence length {seq_len}"
+        )
+
+    batch_sizes: list[int] = []
+    chunks: list[Tensor] = []
+    for t in range(max_len):
+        bs = int((sorted_lengths > t).sum().item())
+        batch_sizes.append(bs)
+        if bs == 0:
+            break
+        chunks.append(input_[t, :bs])
+
+    if not chunks:
+        data = input_[:0]
+    else:
+        data = lucid.concatenate(tuple(chunks), axis=0)
+
+    return PackedSequence(
+        data=data,
+        batch_sizes=Tensor(batch_sizes, device=input_.device),
+        sorted_indices=sorted_indices,
+        unsorted_indices=unsorted_indices,
+    )
+
+
+def pad_packed_sequence(
+    sequence: PackedSequence, batch_first: bool = False, padding_value: _Scalar = 0
+) -> tuple[Tensor, Tensor]:
+    data = sequence.data
+    batch_sizes = sequence.batch_sizes
+    if batch_sizes.ndim != 1:
+        raise ValueError("batch_sizes must be 1D")
+
+    max_len = int(batch_sizes.shape[0])
+    if max_len == 0:
+        raise ValueError("batch_sizes must be non-empty")
+
+    batch_size = int(batch_sizes[0].item())
+    trailing_shape = data.shape[1:]
+
+    if batch_first:
+        out_shape = (batch_size, max_len, *trailing_shape)
+    else:
+        out_shape = (max_len, batch_size, *trailing_shape)
+
+    output = lucid.full(out_shape, padding_value, dtype=data.dtype, device=data.device)
+
+    lengths = [0] * batch_size
+    offset = 0
+    for t in range(max_len):
+        bs = int(batch_sizes[t].item())
+        if bs == 0:
+            break
+
+        chunk = data[offset : offset + bs]
+        offset += bs
+        for i in range(bs):
+            lengths[i] += 1
+        if batch_first:
+            output[:bs, t] = chunk
+        else:
+            output[t, :bs] = chunk
+
+    lengths_t = Tensor(lengths, device=data.device)
+    if sequence.unsorted_indices is not None:
+        if batch_first:
+            output = output[sequence.unsorted_indices]
+        else:
+            output = output[:, sequence.unsorted_indices]
+        lengths_t = lengths_t[sequence.unsorted_indices]
+
+    return output, lengths_t
+
+
+def pack_sequence(
+    sequences: Iterable[Tensor], enforce_sorted: bool = True
+) -> PackedSequence:
+    seq_list = list(sequences)
+    if not seq_list:
+        raise ValueError("pack_sequence expected a non-empty iterable of Tensors")
+
+    lengths = [seq.shape[0] for seq in seq_list]
+    padded = pad_sequence(seq_list, batch_first=False, padding_value=0.0)
+    return pack_padded_sequence(
+        padded, lengths, batch_first=False, enforce_sorted=enforce_sorted
+    )
+
+
+def unpack_sequence(
+    sequence: PackedSequence, batch_first: bool = False
+) -> list[Tensor]:
+    padded, lengths = pad_packed_sequence(
+        sequence, batch_first=batch_first, padding_value=0.0
+    )
+    result: list[Tensor] = []
+    for i, length in enumerate(lengths):
+        l = int(length.item())
+        if batch_first:
+            result.append(padded[i, :l])
+        else:
+            result.append(padded[:l, i])
+    return result

lucid/transforms/image.py

@@ -35,8 +35,8 @@ def add_batch_dim(func: Callable[..., Tensor]) -> Callable:
 class Normalize(nn.Module):
     def __init__(self, mean: tuple[float, ...], std: tuple[float, ...]) -> None:
         super().__init__()
-        self.mean = lucid.tensor(mean)
-        self.std = lucid.tensor(std)
+        self.mean = lucid.tensor(mean).reshape(1, len(mean), 1, 1)
+        self.std = lucid.tensor(std).reshape(1, len(std), 1, 1)
 
     @add_batch_dim
     def forward(self, img: Tensor) -> Tensor:

lucid/visual/__init__.py

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

lucid/visual/mermaid.py

@@ -9,7 +9,7 @@ from lucid._tensor import Tensor
 from lucid.types import _ShapeLike
 
 
-__all__ = ["build_mermaid_chart"]
+__all__ = ["build_tensor_mermaid_chart", "build_module_mermaid_chart"]
 
 
 _NN_MODULES_PREFIX = "lucid.nn.modules."
@@ -255,7 +255,7 @@ def _collapse_repeated_children(
     return out
 
 
-def build_mermaid_chart(
+def build_module_mermaid_chart(
     module: nn.Module,
     input_shape: _ShapeLike | list[_ShapeLike] | None = None,
     inputs: Iterable[Tensor] | Tensor | None = None,
@@ -751,6 +751,192 @@ def build_mermaid_chart(
     return text
 
 
+def build_mermaid_chart(
+    module: nn.Module,
+    input_shape: _ShapeLike | list[_ShapeLike] | None = None,
+    inputs: Iterable[Tensor] | Tensor | None = None,
+    depth: int = 2,
+    direction: str = "LR",
+    include_io: bool = True,
+    show_params: bool = False,
+    return_lines: bool = False,
+    copy_to_clipboard: bool = False,
+    compact: bool = False,
+    use_class_defs: bool = False,
+    end_semicolons: bool = True,
+    edge_mode: Literal["dataflow", "execution"] = "execution",
+    collapse_repeats: bool = True,
+    repeat_min: int = 2,
+    color_by_subpackage: bool = True,
+    container_name_from_attr: bool = True,
+    edge_stroke_width: float = 2.0,
+    emphasize_model_title: bool = True,
+    model_title_font_px: int = 20,
+    show_shapes: bool = False,
+    hide_subpackages: Iterable[str] = (),
+    hide_module_names: Iterable[str] = (),
+    dash_multi_input_edges: bool = True,
+    subgraph_fill: str = "#000000",
+    subgraph_fill_opacity: float = 0.05,
+    subgraph_stroke: str = "#000000",
+    subgraph_stroke_opacity: float = 0.75,
+    force_text_color: str | None = None,
+    edge_curve: str = "natural",
+    node_spacing: int = 50,
+    rank_spacing: int = 50,
+    **forward_kwargs,
+) -> str | list[str]:
+    return build_module_mermaid_chart(
+        module,
+        input_shape=input_shape,
+        inputs=inputs,
+        depth=depth,
+        direction=direction,
+        include_io=include_io,
+        show_params=show_params,
+        return_lines=return_lines,
+        copy_to_clipboard=copy_to_clipboard,
+        compact=compact,
+        use_class_defs=use_class_defs,
+        end_semicolons=end_semicolons,
+        edge_mode=edge_mode,
+        collapse_repeats=collapse_repeats,
+        repeat_min=repeat_min,
+        color_by_subpackage=color_by_subpackage,
+        container_name_from_attr=container_name_from_attr,
+        edge_stroke_width=edge_stroke_width,
+        emphasize_model_title=emphasize_model_title,
+        model_title_font_px=model_title_font_px,
+        show_shapes=show_shapes,
+        hide_subpackages=hide_subpackages,
+        hide_module_names=hide_module_names,
+        dash_multi_input_edges=dash_multi_input_edges,
+        subgraph_fill=subgraph_fill,
+        subgraph_fill_opacity=subgraph_fill_opacity,
+        subgraph_stroke=subgraph_stroke,
+        subgraph_stroke_opacity=subgraph_stroke_opacity,
+        force_text_color=force_text_color,
+        edge_curve=edge_curve,
+        node_spacing=node_spacing,
+        rank_spacing=rank_spacing,
+        **forward_kwargs,
+    )
+
+
+def build_tensor_mermaid_chart(
+    tensor: Tensor,
+    horizontal: bool = False,
+    title: str | None = None,
+    start_id: int | None = None,
+    end_semicolons: bool = True,
+    copy_to_clipboard: bool = False,
+    use_class_defs: bool = True,
+    op_fill: str = "lightgreen",
+    param_fill: str = "plum",
+    result_fill: str = "lightcoral",
+    leaf_fill: str = "lightgray",
+    grad_fill: str = "lightblue",
+    start_fill: str = "gold",
+    stroke_color: str = "#666",
+    stroke_width_px: int = 1,
+) -> str:
+    direction = "LR" if horizontal else "TD"
+    lines: list[str] = [f"flowchart {direction}"]
+    if title:
+        lines.append(f"%% {title}")
+
+    result_id: int = id(tensor)
+    visited: set[int] = set()
+    nodes_to_draw: list[Tensor] = []
+
+    def dfs(t: Tensor) -> None:
+        if id(t) in visited:
+            return
+        visited.add(id(t))
+        for p in t._prev:
+            dfs(p)
+        nodes_to_draw.append(t)
+
+    def tensor_node_id(t: Tensor) -> str:
+        return f"t_{id(t)}"
+
+    def op_node_id(op: object) -> str:
+        return f"op_{id(op)}"
+
+    def add_node(node_id: str, label: str, kind: str) -> None:
+        if node_id in defined_nodes:
+            return
+        defined_nodes.add(node_id)
+        if kind == "op":
+            lines.append(f'{node_id}(("{label}"))')
+        else:
+            lines.append(f'{node_id}["{label}"]')
+
+    dfs(tensor)
+
+    defined_nodes: set[str] = set()
+    edge_lines: list[str] = []
+    class_lines: list[str] = []
+
+    for t in nodes_to_draw:
+        t_id = tensor_node_id(t)
+
+        if not t.is_leaf and t._op is not None:
+            op_id = op_node_id(t._op)
+            op_label = type(t._op).__name__
+            add_node(op_id, op_label, "op")
+            edge_lines.append(f"{op_id} --> {t_id}")
+            class_lines.append(f"class {op_id} op")
+            for inp in t._prev:
+                edge_lines.append(f"{tensor_node_id(inp)} --> {op_id}")
+
+        shape_label = str(t.shape) if t.ndim > 0 else str(t.item())
+        add_node(t_id, shape_label, "tensor")
+
+        if start_id is not None and id(t) == start_id:
+            class_lines.append(f"class {t_id} start")
+        elif isinstance(t, nn.Parameter):
+            class_lines.append(f"class {t_id} param")
+        elif id(t) == result_id:
+            class_lines.append(f"class {t_id} result")
+        elif not t.requires_grad:
+            class_lines.append(f"class {t_id} leaf")
+        else:
+            class_lines.append(f"class {t_id} grad")
+
+    lines.extend(edge_lines)
+    if use_class_defs:
+        lines.append(
+            f"classDef op fill:{op_fill},stroke:{stroke_color},stroke-width:{stroke_width_px}px;"
+        )
+        lines.append(
+            f"classDef param fill:{param_fill},stroke:{stroke_color},stroke-width:{stroke_width_px}px;"
+        )
+        lines.append(
+            f"classDef result fill:{result_fill},stroke:{stroke_color},stroke-width:{stroke_width_px}px;"
+        )
+        lines.append(
+            f"classDef leaf fill:{leaf_fill},stroke:{stroke_color},stroke-width:{stroke_width_px}px;"
+        )
+        lines.append(
+            f"classDef grad fill:{grad_fill},stroke:{stroke_color},stroke-width:{stroke_width_px}px;"
+        )
+        lines.append(
+            f"classDef start fill:{start_fill},stroke:{stroke_color},stroke-width:{stroke_width_px}px;"
+        )
+        lines.extend(class_lines)
+
+    if end_semicolons:
+        lines = [
+            f"{line};" if line and not line.endswith(";") else line for line in lines
+        ]
+
+    text = "\n".join(lines)
+    if copy_to_clipboard:
+        _copy_to_clipboard(text)
+    return text
+
+
 def _copy_to_clipboard(text: str) -> None:
     import os
     import shutil

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lucid-dl
-Version: 2.11.4
+Version: 2.12.0
 Summary: Lumerico's Comprehensive Interface for Deep Learning
 Home-page: https://github.com/ChanLumerico/lucid
 Author: ChanLumerico
@@ -48,35 +48,9 @@ Whether you're a student, educator, or an advanced researcher seeking to demysti
 
 ### 🔥 What's New
 
-- Added additional `nn.Module` hooks for richer introspection during training:
+- Added new visual tool: `lucid.visual.build_tensor_mermaid_chart` which builds a Mermaid chart of given tensor's computatoinal graph
 
-  ```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)
-  ```
+- Added additional `nn.Module` hooks for richer introspection during training:
 
 ## 🔧 How to Install
 

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

@@ -21,16 +21,16 @@ lucid/autograd/__init__.py,sha256=hDoK_B2chRFVhoxsT4vxRKangzBEMWqF8gj2hdoTenk,67
 lucid/data/__init__.py,sha256=qrDIQsnix5ZUEa0yrtomaaWbNJyJ3xEr2gdhRvg70_8,118
 lucid/data/_base.py,sha256=RM8xpBl8qFhm19n7eER_jOsRaxkL3rbOkwUvn6VetSE,5921
 lucid/data/_util.py,sha256=UsbliOrGmM0f1vqppoBPn3RSx53PIqcVx_yVOlHZB6A,1985
-lucid/datasets/__init__.py,sha256=vFlNOP38SSYG75_Yf0Jbyw3okSO7sY4wIjSxIzzIgWg,103
+lucid/datasets/__init__.py,sha256=SY0bCxIUGrSBNfEs4KmOawirmkJ8i-Cyz2CZMBkMsmk,42
 lucid/datasets/_base.py,sha256=yeXPm3La3Beel7U_yPrxjXgGtjndJ3T6NYaQ2_H_Fak,1325
-lucid/datasets/cifar.py,sha256=Wf9r6PSgwFuuYsnZDEQOFNB-06AkqW4DrVsVPRyxEF0,3704
+lucid/datasets/cifar.py,sha256=r8KX-j6svhx3Kk1hSNhexUsJws7Bj31PlLMVRS68dC4,13080
 lucid/datasets/mnist.py,sha256=PUXW2UwmlXJFVJiNkI9Jm58Qe4qWHGA63znkk-y9INM,8603
 lucid/einops/__init__.py,sha256=9Dlmfw6PsIU9b_a89Zre4yV2rztRHPCL4QpsUnXJwjM,802
 lucid/einops/_func.py,sha256=XXsX9lse_0turKoFnOTtLdY6hBUi0gq_8K81G7nr80I,21026
 lucid/linalg/__init__.py,sha256=N-LrlC3qSsOMt6Ad1-PP3Qc3QH6EWNf5P50GBvwb9aQ,1118
 lucid/linalg/_func.py,sha256=Iyeut5nHwQmO8N326kQUaTjgoKVoBaxt_gy_3NXXD60,16378
 lucid/models/__init__.py,sha256=wegfOBvwJTFFee8eVt90zJoLsbbEpdT5G2y-mpO5xcE,89
-lucid/models/util.py,sha256=2g8FLcMLRgVxgGEaYuwJyFxeXu-A_a4_MVr0K-TNh74,5195
+lucid/models/utils.py,sha256=2g8FLcMLRgVxgGEaYuwJyFxeXu-A_a4_MVr0K-TNh74,5195
 lucid/models/imgclf/__init__.py,sha256=kQH-nNu8_TPJ7Av151WSpcY4GJ06gGAd6Ozs3m3KMcE,590
 lucid/models/imgclf/alex.py,sha256=fZsPdCjWUseCrxBwKj-i5fPSDYLgBpfm0SJe07YKRuE,1472
 lucid/models/imgclf/coatnet.py,sha256=HKjpy-lBKgz743EijT7jEeMxYjrZHzgU5fOrgtZfxYg,13720
@@ -55,7 +55,7 @@ lucid/models/imgclf/senet.py,sha256=I5o9eHWzquNyLqZM4thMtZtIBDYGczjARl1Isx6GyCk,
 lucid/models/imgclf/sknet.py,sha256=rENInsSB2yLXJ7A9kWZ-9lDFXcKaUOIpzV0359umPRI,4535
 lucid/models/imgclf/swin.py,sha256=lClJTX6ObF1PuzYR99Grgc7AhignbomwYFvqkQoCMx4,27969
 lucid/models/imgclf/vgg.py,sha256=fWy78AAHJre3Msy4DK5nhQwThI-7frsdqRS-JYtFiXM,2457
-lucid/models/imgclf/vit.py,sha256=NXzPIiyXxcE1-g25m36-_YwKnJZ0gl1-jf7G3V12jS0,3594
+lucid/models/imgclf/vit.py,sha256=AUwsueQh9PY9d5org1PQjYzjSs9TVDOYElO9daO9Za8,3656
 lucid/models/imgclf/xception.py,sha256=Y7YKCzF_y4r864hLouW0eE7M-kxA59SiI3-iIFsXVhQ,3728
 lucid/models/imgclf/zfnet.py,sha256=brH5tHLVWTUfCqu-BwfFb0yZV9p5DmXN4O6cyP3U26U,1469
 lucid/models/imggen/__init__.py,sha256=J6MlEHqXxAYINbeQmyb85ev_IEOvQDTxTQjPgX6hdpY,59
@@ -76,11 +76,10 @@ lucid/models/objdet/yolo/yolo_v3.py,sha256=B5U42Npwfg8nSgU9E261zf0cbQS9RVYrX1ADD
 lucid/models/objdet/yolo/yolo_v4.py,sha256=RFbBumreXmy6s8IYZvUuhW0893ss8sx_8Vgi6KbBKWo,21467
 lucid/models/seq2seq/__init__.py,sha256=wjsrhj4H_AcqwwbebAN8b68QBA8L6p1_12dkG2995-w,27
 lucid/models/seq2seq/transformer.py,sha256=y5rerCs1s6jXTsVvbgscWScKpQKuSu1fezsBe7PNTRA,3513
-lucid/nn/__init__.py,sha256=_hk6KltQIJuWXowXstMSu3TjiaTP8zMLNvGpjnA9Mpw,182
+lucid/nn/__init__.py,sha256=nyy6px1CxfchWUh68xCiQSxD7Gk65vamhWK8ztRvH68,184
 lucid/nn/fused.py,sha256=75fcXuo6fHSO-JtjuKhowhHSDr4qc5871WR63sUzH0g,5492
 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
 lucid/nn/_kernel/activation.py,sha256=mfe48Aw3_Hv0hZEVC7DxDw19XK9XSLfdCOvo2JcZz_o,5662
 lucid/nn/_kernel/attention.py,sha256=1k0gboLObMNVow2v3TwliXC_2v8uKf2o8jHYFuyQqcg,3699
@@ -112,10 +111,13 @@ lucid/nn/modules/linear.py,sha256=87cuFWYct9JlmtVC3jGR-8eouxxzANaVA6cd7p9r2Ho,28
 lucid/nn/modules/loss.py,sha256=pjEMIruhtpTHhHFsNThS9LFz-aI_DAXLqMV8KRXydEg,3431
 lucid/nn/modules/norm.py,sha256=bYsKOg58kxzhMhbyvHrDDgVzN_p3D9HBTdYWpDtDeHQ,6842
 lucid/nn/modules/pool.py,sha256=ymVnS2NZjh08Tw0VeOfkB6AVrMeLmCKvgxkmEO3KUuw,5044
-lucid/nn/modules/rnn.py,sha256=lsvQZiEHm1wGbiNWKQngAle7MbqGaXSBM1LUieCaZIk,17233
+lucid/nn/modules/rnn.py,sha256=L2rqFRcdr0U33YFeVvthDwDFIE98PrO-OjFiX9IzlIs,21098
 lucid/nn/modules/sparse.py,sha256=EpjiviED2nI55wUjh1twFwa4Lvlrzw0TR6lpCDGeSbo,1147
 lucid/nn/modules/transformer.py,sha256=z56emF_eX18pxRELjfmmsY-7Bn9h2yjIdxCaxs6YDwA,11246
 lucid/nn/modules/vision.py,sha256=8xYasT7TNj4NXwMwwJIw1nbV1paeWEFg_ZohXn9kZBg,1579
+lucid/nn/utils/__init__.py,sha256=ynHrPi9SPdRRXhGjghG42FRBcEiVN8Hb_04XHBZqy_o,46
+lucid/nn/utils/_grad.py,sha256=8EFN7TDHb09LHXK9dPjAdSLgGnL3r48Ct2rYztXKQxM,2335
+lucid/nn/utils/rnn.py,sha256=yJIktD-cbFvegzyDrif4aQFshpF64cCxAweCikrKm7s,6963
 lucid/optim/__init__.py,sha256=21EcCCPwrhPGP9TXvDje075_S2hPr0pHToygCaq8keI,201
 lucid/optim/_base.py,sha256=KxM5h5ONeO8hCpAzD2_vverFRKeymu2XC6AHN_L_v3g,4859
 lucid/optim/ada.py,sha256=-WQcC81oSYw3ffa59dPuNtDZfJ1KDrUw3zyKuPn5h5Y,6451
@@ -129,14 +131,13 @@ lucid/random/__init__.py,sha256=s8EAaKhEiTKT_vYjP4IFHx0xQVa1jqc_qIyvMauUu7M,2727
 lucid/random/_func.py,sha256=1Lu4m-ciEK037chNDGqv_j00RgGGzQ7UfslSfYActUk,2232
 lucid/transforms/__init__.py,sha256=DGznMbqhXdU9FLDMKnJawScO4HCqu40Sf_j4vJGJrjc,90
 lucid/transforms/_base.py,sha256=v3elm7l0VoWvrT_qgoJiRzLH42tHoUcPIKNaPuxI_2E,1448
-lucid/transforms/image.py,sha256=S9gZzMck4EQSmDQZ3ATi2fsUh4-hqFqeDjhMMJe8TdU,3762
-lucid/visual/__init__.py,sha256=NfHhHYNVv9mQQ4MST3-OAIkAcFyYrihJC4qUf88DySI,44
-lucid/visual/graph.py,sha256=ZSlrJI3dQwYjz8XbgAfNd8-8YuH9Ji7Mz1J6UsnHTaI,4711
-lucid/visual/mermaid.py,sha256=87hFe4l9EYP6Cg2l2hP2INQiBHKkgVClH5nBWFY9ddY,26499
+lucid/transforms/image.py,sha256=Pn4AFQ5nQixOLmlpiSlxVd8tyALOvg24UgDueY-U8fc,3817
+lucid/visual/__init__.py,sha256=tRgyNHzKWA8cp-a_GV586Bs0yJUN5ZTmKgnUhscutHQ,23
+lucid/visual/mermaid.py,sha256=m0X0kkdLuCxEzKmXSy3zplUaa3Gov8RRonKyHiEvfHE,32738
 lucid/weights/__init__.py,sha256=z1AikA3rOEeckWGkYWlcZkxNlJo9Xwa39PL6ly3hWnc,8801
 lucid/weights/__init__.pyi,sha256=lFonYC3cUx2Idolf3AEPnjFcyqcn3UDU84oJlZafqLY,3013
-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,,
+lucid_dl-2.12.0.dist-info/licenses/LICENSE,sha256=vxRFYnVD1IeYtsvw-KmoElfqrjxKHv1h9YTvsG54loQ,1065
+lucid_dl-2.12.0.dist-info/METADATA,sha256=Y7doYNmgXQugwLzkYsJBv4Jzw1g9ZMsIxXYofaCmdAc,11679
+lucid_dl-2.12.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+lucid_dl-2.12.0.dist-info/top_level.txt,sha256=uzP_qBx9iNWIHKJRlElYcBLYVqMpdm9Q1Ma63QPYbFc,6
+lucid_dl-2.12.0.dist-info/RECORD,,

lucid/visual/graph.py

@@ -1,141 +0,0 @@
-from typing import Union
-from warnings import deprecated
-
-import networkx as nx
-import matplotlib.pyplot as plt
-
-import lucid.nn as nn
-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,
-    title: Union[str, None] = None,
-    start_id: Union[int, None] = None,
-) -> plt.Figure:
-    G: nx.DiGraph = nx.DiGraph()
-    result_id: int = id(tensor)
-
-    visited: set[int] = set()
-    nodes_to_draw: list[Tensor] = []
-
-    def dfs(t: Tensor) -> None:
-        if id(t) in visited:
-            return
-        visited.add(id(t))
-        for p in t._prev:
-            dfs(p)
-        nodes_to_draw.append(t)
-
-    dfs(tensor)
-
-    for t in nodes_to_draw:
-        if not t.is_leaf and t._op is not None:
-            op_id: int = id(t._op)
-            op_label: str = type(t._op).__name__
-            G.add_node(op_id, label=op_label, shape="circle", color="lightgreen")
-            G.add_edge(op_id, id(t))
-            for inp in t._prev:
-                G.add_edge(id(inp), op_id)
-
-        shape_label: str = str(t.shape) if t.ndim > 0 else str(t.item())
-        if isinstance(t, nn.Parameter):
-            color: str = "plum"
-        else:
-            color = (
-                "lightcoral"
-                if id(t) == result_id
-                else "lightgray" if not t.requires_grad else "lightblue"
-            )
-        if start_id is not None and id(t) == start_id:
-            color = "gold"
-
-        G.add_node(id(t), label=shape_label, shape="rectangle", color=color)
-
-    def grid_layout(
-        G: nx.DiGraph, horizontal: bool = False
-    ) -> tuple[dict, tuple, float, int]:
-        levels: dict[int, int] = {}
-        for node in nx.topological_sort(G):
-            preds = list(G.predecessors(node))
-            levels[node] = 0 if not preds else max(levels[p] for p in preds) + 1
-
-        level_nodes: dict[int, list[int]] = {}
-        for node, level in levels.items():
-            level_nodes.setdefault(level, []).append(node)
-
-        def autoscale(
-            level_nodes: dict[int, list[int]],
-            horizontal: bool = False,
-            base_size: float = 0.5,
-            base_nodesize: int = 500,
-        ) -> tuple[tuple[float, float], float, int]:
-            num_levels: int = len(level_nodes)
-            max_width: int = max(len(nodes) for nodes in level_nodes.values())
-            node_count: int = sum(len(nodes) for nodes in level_nodes.values())
-
-            if horizontal:
-                fig_w: float = min(32, max(4.0, base_size * num_levels))
-                fig_h: float = min(32, max(4.0, base_size * max_width))
-            else:
-                fig_w = min(32, max(4.0, base_size * max_width))
-                fig_h = min(32, max(4.0, base_size * num_levels))
-
-            nodesize: float = (
-                base_nodesize
-                if node_count <= 100
-                else base_nodesize * (100 / node_count)
-            )
-            fontsize: int = max(5, min(8, int(80 / node_count)))
-            return (fig_w, fig_h), nodesize, fontsize
-
-        figsize, nodesize, fontsize = autoscale(level_nodes, horizontal)
-        pos: dict[int, tuple[float, float]] = {}
-        for level, nodes in level_nodes.items():
-            for i, node in enumerate(nodes):
-                pos[node] = (
-                    (level * 2.5, -i * 2.0) if horizontal else (i * 2.5, -level * 2.0)
-                )
-        return pos, figsize, nodesize, fontsize
-
-    labels: dict[int, str] = nx.get_node_attributes(G, "label")
-    colors: dict[int, str] = nx.get_node_attributes(G, "color")
-    shapes: dict[int, str] = nx.get_node_attributes(G, "shape")
-    pos, figsize, nodesize, fontsize = grid_layout(G, horizontal)
-
-    fig, ax = plt.subplots(figsize=figsize)
-
-    rect_nodes: list[int] = [n for n in G.nodes() if shapes.get(n) == "rectangle"]
-    circ_nodes: list[int] = [n for n in G.nodes() if shapes.get(n) == "circle"]
-    rect_colors: list[str] = [colors[n] for n in rect_nodes]
-
-    nx.draw_networkx_nodes(
-        G,
-        pos,
-        nodelist=rect_nodes,
-        node_color=rect_colors,
-        node_size=nodesize,
-        node_shape="s",
-        ax=ax,
-    )
-    nx.draw_networkx_nodes(
-        G,
-        pos,
-        nodelist=circ_nodes,
-        node_color="lightgreen",
-        node_size=nodesize,
-        node_shape="o",
-        ax=ax,
-    )
-    nx.draw_networkx_edges(G, pos, width=0.5, arrows=True, edge_color="gray", ax=ax)
-    nx.draw_networkx_labels(G, pos, labels=labels, font_size=fontsize, ax=ax)
-
-    ax.axis("off")
-    ax.set_title(title if title is not None else "")
-
-    return fig