lucid-dl 2.11.5__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_dl-2.11.5.dist-info → lucid_dl-2.12.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lucid-dl
-Version: 2.11.5
+Version: 2.12.0
 Summary: Lumerico's Comprehensive Interface for Deep Learning
 Home-page: https://github.com/ChanLumerico/lucid
 Author: ChanLumerico

{lucid_dl-2.11.5.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,13 +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/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.5.dist-info/licenses/LICENSE,sha256=vxRFYnVD1IeYtsvw-KmoElfqrjxKHv1h9YTvsG54loQ,1065
-lucid_dl-2.11.5.dist-info/METADATA,sha256=kEkdiWHxh_TEYmZX3lr_awR_02Ef87QqSwtj4j98T-o,11679
-lucid_dl-2.11.5.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-lucid_dl-2.11.5.dist-info/top_level.txt,sha256=uzP_qBx9iNWIHKJRlElYcBLYVqMpdm9Q1Ma63QPYbFc,6
-lucid_dl-2.11.5.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,,