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

lucid/__init__.py

@@ -25,7 +25,7 @@ import json
 import math
 import numpy as np
 
-from lucid._tensor import Tensor
+from lucid._tensor import *
 from lucid._func import *
 from lucid._util import *
 
@@ -308,7 +308,7 @@ def register_model(func: _ModuleReturnFunc) -> _ModuleReturnFunc:
 
 
 def _conv_view_limit_mb() -> int:
-    from lucid._kernel import conv as _conv_kernel
+    from lucid.nn._kernel import conv as _conv_kernel
 
     return _conv_kernel.get_conv_view_limit_mb()
 

lucid/_tensor/__init__.py

@@ -1 +1,11 @@
-from lucid._tensor.tensor import Tensor
+from lucid._tensor.tensor import (
+    Tensor,
+    LongTensor,
+    IntTensor,
+    ShortTensor,
+    CharTensor,
+    HalfTensor,
+    FloatTensor,
+    DoubleTensor,
+    BoolTensor,
+)

lucid/_tensor/base.py

@@ -108,6 +108,8 @@ class _TensorBase:
 
     def broadcast_to(self, shape: _ShapeLike) -> Self: ...
 
+    def expand(self, *sizes: int | _ShapeLike) -> Self: ...
+
     def chunk(self, chunks: int, axis: int = 0) -> tuple[Self, ...]: ...
 
     def swapaxes(self, axis1: int, axis2: int) -> Self: ...

lucid/_tensor/tensor.py

@@ -1,4 +1,15 @@
-from typing import Callable, Iterator, Optional, Self, SupportsIndex, Any, overload
+from typing import (
+    Callable,
+    Iterator,
+    Optional,
+    Self,
+    SupportsIndex,
+    Any,
+    overload,
+    Generic,
+    TypeVar,
+    ClassVar,
+)
 from types import NoneType
 from collections import deque
 
@@ -22,15 +33,32 @@ from lucid._backend.core import BackwardOperation, Operation, noop
 from lucid._backend.metal import mx, parse_mlx_indexing, check_metal_availability
 
 
+__all__ = [
+    "Tensor",
+    "FloatTensor",
+    "DoubleTensor",
+    "HalfTensor",
+    "CharTensor",
+    "ShortTensor",
+    "IntTensor",
+    "LongTensor",
+    "BoolTensor",
+]
+
+
+DType = TypeVar("DType", bound=Numeric | bool)
+
 _HookType = Callable[["Tensor", _NumPyArray | _MLXArray], None]
 
 _dtype_map = {int: types.Int64, float: types.Float64, complex: types.Complex64}
 
 
-class Tensor(_TensorBase, _TensorInplace):
+class Tensor(Generic[DType], _TensorBase, _TensorInplace):
+    _fixed_dtype: ClassVar[Numeric | None] = None
+
     def __init__(
         self,
-        data: _ArrayOrScalar | _MLXArray,
+        data: _ArrayOrScalar,
         requires_grad: bool = False,
         keep_grad: bool = False,
         dtype: _BuiltinNumeric | Numeric | None = None,
@@ -39,6 +67,9 @@ class Tensor(_TensorBase, _TensorInplace):
         self._is_free = False
         self._is_bool_tensor = False
 
+        if self._fixed_dtype is not None:
+            dtype = self._fixed_dtype
+
         if dtype is bool:
             self._is_bool_tensor = True
             dtype = None
@@ -285,6 +316,12 @@ class Tensor(_TensorBase, _TensorInplace):
             dtype = device_or_dtype
             return self.astype(dtype)
 
+    def cpu(self) -> Self:
+        return self.to(device="cpu")
+
+    def gpu(self) -> Self:
+        return self.to(device="gpu")
+
     def is_cpu(self) -> bool:
         return self.device == "cpu"
 
@@ -480,3 +517,155 @@ class Tensor(_TensorBase, _TensorInplace):
 
     def bool(self) -> Self:
         return self.astype(bool)
+
+
+class LongTensor(Tensor[types.Int64]):
+    _fixed_dtype: ClassVar[Numeric | None] = types.Int64
+
+    def __init__(
+        self,
+        data: _ArrayOrScalar,
+        requires_grad: bool = False,
+        keep_grad: bool = False,
+        device: _DeviceType = "cpu",
+    ) -> None:
+        super().__init__(
+            data=data,
+            requires_grad=requires_grad,
+            keep_grad=keep_grad,
+            dtype=types.Int64,
+            device=device,
+        )
+
+
+class IntTensor(Tensor[types.Int32]):
+    _fixed_dtype: ClassVar[Numeric | None] = types.Int32
+
+    def __init__(
+        self,
+        data: _ArrayOrScalar,
+        requires_grad: bool = False,
+        keep_grad: bool = False,
+        device: _DeviceType = "cpu",
+    ) -> None:
+        super().__init__(
+            data=data,
+            requires_grad=requires_grad,
+            keep_grad=keep_grad,
+            dtype=types.Int32,
+            device=device,
+        )
+
+
+class ShortTensor(Tensor[types.Int16]):
+    _fixed_dtype: ClassVar[Numeric | None] = types.Int16
+
+    def __init__(
+        self,
+        data: _ArrayOrScalar,
+        requires_grad: bool = False,
+        keep_grad: bool = False,
+        device: _DeviceType = "cpu",
+    ) -> None:
+        super().__init__(
+            data=data,
+            requires_grad=requires_grad,
+            keep_grad=keep_grad,
+            dtype=types.Int16,
+            device=device,
+        )
+
+
+class CharTensor(Tensor[types.Int8]):
+    _fixed_dtype: ClassVar[Numeric | None] = types.Int8
+
+    def __init__(
+        self,
+        data: _ArrayOrScalar,
+        requires_grad: bool = False,
+        keep_grad: bool = False,
+        device: _DeviceType = "cpu",
+    ) -> None:
+        super().__init__(
+            data=data,
+            requires_grad=requires_grad,
+            keep_grad=keep_grad,
+            dtype=types.Int8,
+            device=device,
+        )
+
+
+class HalfTensor(Tensor[types.Float16]):
+    _fixed_dtype: ClassVar[Numeric | None] = types.Float16
+
+    def __init__(
+        self,
+        data: _ArrayOrScalar,
+        requires_grad: bool = False,
+        keep_grad: bool = False,
+        device: _DeviceType = "cpu",
+    ) -> None:
+        super().__init__(
+            data=data,
+            requires_grad=requires_grad,
+            keep_grad=keep_grad,
+            dtype=types.Float16,
+            device=device,
+        )
+
+
+class FloatTensor(Tensor[types.Float32]):
+    _fixed_dtype: ClassVar[Numeric | None] = types.Float32
+
+    def __init__(
+        self,
+        data: _ArrayOrScalar,
+        requires_grad: bool = False,
+        keep_grad: bool = False,
+        device: _DeviceType = "cpu",
+    ) -> None:
+        super().__init__(
+            data=data,
+            requires_grad=requires_grad,
+            keep_grad=keep_grad,
+            dtype=types.Float32,
+            device=device,
+        )
+
+
+class DoubleTensor(Tensor[types.Float64]):
+    _fixed_dtype: ClassVar[Numeric | None] = types.Float64
+
+    def __init__(
+        self,
+        data: _ArrayOrScalar,
+        requires_grad: bool = False,
+        keep_grad: bool = False,
+        device: _DeviceType = "cpu",
+    ) -> None:
+        super().__init__(
+            data=data,
+            requires_grad=requires_grad,
+            keep_grad=keep_grad,
+            dtype=types.Float64,
+            device=device,
+        )
+
+
+class BoolTensor(Tensor[bool]):
+    _fixed_dtype: ClassVar[Numeric | None] = None
+
+    def __init__(
+        self,
+        data: _ArrayOrScalar,
+        requires_grad: bool = False,
+        keep_grad: bool = False,
+        device: _DeviceType = "cpu",
+    ) -> None:
+        super().__init__(
+            data=data,
+            requires_grad=requires_grad,
+            keep_grad=keep_grad,
+            dtype=bool,
+            device=device,
+        )

lucid/_util/__init__.py

@@ -9,11 +9,11 @@ from lucid._util import func
 # fmt: off
 __all__ = [
     "reshape", "squeeze", "unsqueeze", "expand_dims", "ravel", "stack", "hstack",
-    "vstack", "concatenate", "pad", "repeat", "tile", "flatten", "meshgrid", 
-    "split", "tril", "triu", "broadcast_to", "chunk", "masked_fill", "roll", 
-    "unbind", "sort", "nonzero", "unique", "topk", "argsort", "histogramdd", 
-    "histogram", "histogram2d", "where", "nonzero", "argmin", "argmax", 
-    "diagonal",
+    "vstack", "concatenate", "pad", "repeat", "tile", "flatten", "meshgrid",
+    "split", "tril", "triu", "broadcast_to", "expand", "chunk", "masked_fill",
+    "roll", "unbind", "sort", "nonzero", "unique", "topk", "argsort",
+    "histogramdd", "histogram", "histogram2d", "where", "nonzero", "argmin",
+    "argmax", "diagonal",
 ]
 # fmt: on
 
@@ -106,6 +106,14 @@ def broadcast_to(a: Tensor, /, shape: _ShapeLike) -> Tensor:
     return func.broadcast_to(shape)(a)
 
 
+def expand(a: Tensor, /, *sizes: int | _ShapeLike) -> Tensor:
+    if len(sizes) == 1 and isinstance(sizes[0], (tuple, list)):
+        shape = sizes[0]
+    else:
+        shape = sizes
+    return func.expand(shape)(a)
+
+
 def chunk(a: Tensor, /, chunks: int, axis: int = 0) -> tuple[Tensor, ...]:
     return func.chunk(chunks, axis)(a)
 
@@ -257,6 +265,7 @@ Tensor.split = split
 Tensor.tril = tril
 Tensor.triu = triu
 Tensor.broadcast_to = broadcast_to
+Tensor.expand = expand
 Tensor.chunk = chunk
 Tensor.masked_fill = masked_fill
 Tensor.roll = roll

lucid/_util/func.py

@@ -605,6 +605,79 @@ class broadcast_to(Operation):
         return self.result.grad.reshape(self.original_shape)
 
 
+class expand(Operation):
+    def __init__(self, shape: _ShapeLike) -> None:
+        super().__init__()
+        self.shape = shape
+
+    def _resolve_shape(self, input_shape: tuple[int, ...]) -> tuple[int, ...]:
+        shape = tuple(int(dim) for dim in self.shape)
+        if len(shape) == 0:
+            raise ValueError("expand() expects at least one dimension.")
+
+        if len(shape) < len(input_shape):
+            raise ValueError(
+                "expand() cannot shrink the number of dimensions from "
+                f"{len(input_shape)} to {len(shape)}."
+            )
+
+        ndim_diff = len(shape) - len(input_shape)
+        padded_input = (1,) * ndim_diff + input_shape
+
+        resolved: list[int] = []
+        for axis, (target_dim, input_dim) in enumerate(zip(shape, padded_input)):
+            if target_dim == -1:
+                if axis < ndim_diff:
+                    raise ValueError(
+                        "expand() cannot use -1 in a leading, "
+                        "non-existing dimension."
+                    )
+                target_dim = input_dim
+
+            elif target_dim < -1:
+                raise ValueError("expand() size must be >= -1.")
+
+            if input_dim == target_dim:
+                resolved.append(target_dim)
+            elif input_dim == 1 and target_dim >= 0:
+                resolved.append(target_dim)
+            else:
+                raise ValueError(
+                    "expand() cannot expand dimension "
+                    f"{axis} from {input_dim} to {target_dim}."
+                )
+
+        return tuple(resolved)
+
+    @unary_func_op()
+    def cpu(self, a: Tensor) -> _FuncOpReturnType:
+        self.original_shape = a.shape
+        self.expanded_shape = self._resolve_shape(a.shape)
+
+        self.result = Tensor(np.broadcast_to(a.data, self.expanded_shape))
+        return self.result, self.__grad__
+
+    @unary_func_op(device="gpu")
+    def gpu(self, a: Tensor) -> _FuncOpReturnType:
+        self.original_shape = a.shape
+        self.expanded_shape = self._resolve_shape(a.shape)
+
+        self.result = Tensor(mx.broadcast_to(a.data, self.expanded_shape))
+        return self.result, self.__grad__
+
+    def __grad__(self) -> _GradType:
+        input_shape = self.original_shape
+        ndim_diff = len(self.expanded_shape) - len(input_shape)
+        if ndim_diff > 0:
+            input_shape = (1,) * ndim_diff + input_shape
+
+        for axis, (in_dim, out_dim) in enumerate(zip(input_shape, self.expanded_shape)):
+            if in_dim == 1 and out_dim > 1:
+                self.result.grad = self.result.grad.sum(axis=axis, keepdims=True)
+
+        return self.result.grad.reshape(self.original_shape)
+
+
 class chunk(Operation):
     def __init__(self, chunks: int, axis: int) -> None:
         super().__init__()

lucid/models/__init__.py

@@ -2,3 +2,4 @@ from .imgclf import *
 from .imggen import *
 from .objdet import *
 from .seq2seq import *
+from .seqclf import *

lucid/models/seqclf/__init__.py

@@ -0,0 +1 @@
+from .bert import *

lucid/models/seqclf/bert.py

@@ -0,0 +1,31 @@
+import lucid
+import lucid.nn as nn
+import lucid.nn.functional as F
+
+from lucid._tensor import Tensor
+
+
+class _BertEmbeddings(nn.Module):
+    def __init__(
+        self,
+        vocab_size: int,
+        hidden_size: int,
+        pad_token_id: int,
+        max_position_embeddings: int,
+        type_vocab_size: int,
+        layer_norm_eps: float,
+        hidden_dropout_prob: float,
+    ) -> None:
+        super().__init__()
+        self.word_embeddings = nn.Embedding(vocab_size, hidden_size, pad_token_id)
+        self.position_embeddings = nn.Embedding(max_position_embeddings, hidden_size)
+        self.token_type_embeddings = nn.Embedding(type_vocab_size)
+
+        self.layernorm = nn.LayerNorm(hidden_size, eps=layer_norm_eps)
+        self.dropout = nn.Dropout(hidden_dropout_prob)
+
+        self.position_ids: nn.Buffer
+        self.register_buffer(
+            "position_ids", nn.Buffer(lucid.arange(max_position_embeddings))
+        )
+        # TODO: Implement `lucid.Tensor.expand`

lucid/nn/_kernel/embedding.py

@@ -1,4 +1,4 @@
-import functools
+from functools import partial
 from types import ModuleType
 
 import numpy as np
@@ -7,49 +7,44 @@ from lucid._backend.core import Operation, func_op, _FuncOpReturnType, _GradType
 from lucid._backend.metal import mx
 from lucid._tensor import Tensor
 
-from lucid.types import _DeviceType, _TensorData
-
-
-def _as_int_array(arr, lib_: ModuleType) -> _TensorData:
-    if lib_ is np:
-        return arr.astype(np.int64)
-    return arr.astype(mx.int32)
-
 
 class embedding_kernel(Operation):
-    def __init__(self) -> None:
+    def __init__(self, padding_idx: int = -1) -> None:
         super().__init__()
+        self.padding_idx = int(padding_idx)
         self._indices = None
         self._num_embeddings = None
 
     def clear(self) -> None:
         super().clear()
+        self.padding_idx = -1
         self._indices = None
         self._num_embeddings = None
 
     @func_op(n_in=2, n_ret=1)
     def cpu(self, indices: Tensor, weight: Tensor) -> _FuncOpReturnType:
-        return self._forward(indices, weight, lib_=np, device="cpu")
+        return self._forward(indices, weight, lib_=np)
 
     @func_op(n_in=2, n_ret=1, device="gpu")
     def gpu(self, indices: Tensor, weight: Tensor) -> _FuncOpReturnType:
-        return self._forward(indices, weight, lib_=mx, device="gpu")
+        return self._forward(indices, weight, lib_=mx)
 
     def _forward(
-        self, indices: Tensor, weight: Tensor, lib_: ModuleType, device: _DeviceType
+        self, indices: Tensor, weight: Tensor, lib_: ModuleType
     ) -> _FuncOpReturnType:
-        idx = _as_int_array(indices.data, lib_)
+        idx = indices.data
         out = weight.data[idx]
 
         self._indices = idx
         self._num_embeddings = int(weight.shape[0])
 
-        self.result = Tensor(out, device=device)
-        return self.result, functools.partial(self.__grad__, lib_=lib_)
+        self.result = Tensor(out)
+        return self.result, partial(self.__grad__, lib_=lib_)
 
     def __grad__(self, lib_: ModuleType) -> _GradType:
         if self.result is None or self.result.grad is None:
             raise RuntimeError("embedding backward called before forward.")
+
         if self._indices is None or self._num_embeddings is None:
             raise RuntimeError("embedding cached data missing.")
 
@@ -58,15 +53,23 @@ class embedding_kernel(Operation):
         grad_flat = grad_out.reshape(idx.shape[0], -1)
 
         if lib_ is np:
+            if self.padding_idx >= 0:
+                keep = idx != self.padding_idx
+                idx = idx[keep]
+                grad_flat = grad_flat[keep]
+
             grad_w = np.zeros(
                 (self._num_embeddings, grad_flat.shape[1]), dtype=grad_out.dtype
             )
             np.add.at(grad_w, idx, grad_flat)
+
         else:
             grad_w = mx.zeros(
                 (self._num_embeddings, grad_flat.shape[1]), dtype=grad_out.dtype
             )
             for i in range(idx.shape[0]):
+                if self.padding_idx >= 0 and int(idx[i]) == self.padding_idx:
+                    continue
                 grad_w = grad_w.at[idx[i]].add(grad_flat[i])
 
         return None, grad_w

lucid/nn/functional/_util.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 import lucid
 import lucid.nn.functional
 
@@ -5,6 +7,7 @@ from lucid._tensor import Tensor
 from lucid.types import _Scalar, Numeric
 
 from lucid.nn._kernel.embedding import embedding_kernel
+from lucid._backend.metal import mx
 
 
 def _interpolate_bilinear(
@@ -131,17 +134,46 @@ def embedding(
     max_norm: float | None = None,
     norm_type: float = 2.0,
 ) -> Tensor:
+    num_embeddings = int(weight.shape[0])
+    if padding_idx is None:
+        pad = -1
+    else:
+        pad = int(padding_idx)
+        if pad < 0:
+            pad += num_embeddings
+        if pad < 0 or pad >= num_embeddings:
+            raise IndexError("padding_idx out of range.")
+
     indices = input_.astype(lucid.Int)
-    op = embedding_kernel()
-    output = op(indices, weight)
-    if padding_idx is not None:
-        mask = input_.data == padding_idx
-        output *= 1 - mask[..., None]
+    idx_data = indices.data
+
+    if (idx_data < 0).any() or (idx_data >= num_embeddings).any():
+        raise IndexError("embedding indices out of range.")
 
     if max_norm is not None:
-        norm = (output**norm_type).sum(axis=-1, keepdims=True) ** (1 / norm_type)
-        scaling = max_norm / (norm + (norm == 0))
-        output *= scaling
+        lib_ = np if weight.is_cpu() else mx
+        flat = idx_data.reshape(-1)
+
+        w = weight.data[flat]
+        if norm_type <= 0:
+            raise ValueError("norm_type must be positive.")
+
+        norms = (lib_.abs(w) ** norm_type).sum(axis=1) ** (1.0 / norm_type)
+        scale = lib_.minimum(1.0, max_norm / (norms + (norms == 0)))
+
+        if pad >= 0:
+            mask = flat == pad
+            mask_f = mask.astype(scale.dtype)
+            scale = scale * (1 - mask_f) + mask_f
+
+        weight.data[flat] = w * scale[:, None]
+
+    op = embedding_kernel(padding_idx=pad)
+    output = op(indices, weight)
+
+    if pad >= 0:
+        mask = input_.data == pad
+        output *= 1 - mask[..., None]
 
     return output
 

lucid/nn/modules/attention.py

@@ -1,3 +1,5 @@
+import math
+
 import lucid
 import lucid.nn as nn
 import lucid.nn.functional as F
@@ -41,6 +43,7 @@ class ScaledDotProductAttention(nn.Module):
     "num_heads",
     "dropout",
     "bias",
+    "use_separate_proj_weight",
     "add_bias_kv",
     "add_zero_attn",
 )
@@ -51,6 +54,7 @@ class MultiHeadAttention(nn.Module):
         num_heads: int,
         dropout: float = 0.0,
         bias: bool = True,
+        use_separate_proj_weight: bool = True,
         add_bias_kv: bool = False,
         add_zero_attn: bool = False,
         kdim: int | None = None,
@@ -60,6 +64,7 @@ class MultiHeadAttention(nn.Module):
         self.embed_dim = embed_dim
         self.num_heads = num_heads
         self.dropout = dropout
+        self.use_separate_proj_weight = use_separate_proj_weight
         self.add_bias_kv = add_bias_kv
         self.add_zero_attn = add_zero_attn
 
@@ -70,9 +75,30 @@ class MultiHeadAttention(nn.Module):
         kdim = kdim if kdim is not None else embed_dim
         vdim = vdim if vdim is not None else embed_dim
 
-        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
-        self.k_proj = nn.Linear(kdim, embed_dim, bias=bias)
-        self.v_proj = nn.Linear(vdim, embed_dim, bias=bias)
+        if use_separate_proj_weight:
+            self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+            self.k_proj = nn.Linear(kdim, embed_dim, bias=bias)
+            self.v_proj = nn.Linear(vdim, embed_dim, bias=bias)
+
+            self.register_parameter("in_proj_weight", None)
+            self.register_parameter("in_proj_bias", None)
+        else:
+            if kdim != embed_dim or vdim != embed_dim:
+                raise ValueError(
+                    "in_proj_weight requires kdim and vdim to equal embed_dim."
+                )
+
+            weight_ = lucid.empty(3 * embed_dim, embed_dim)
+            self.in_proj_weight = nn.Parameter(weight_)
+            if bias:
+                bias_ = lucid.empty(3 * embed_dim)
+                self.in_proj_bias = nn.Parameter(bias_)
+            else:
+                self.register_parameter("in_proj_bias", None)
+
+            self.q_proj = None
+            self.k_proj = None
+            self.v_proj = None
 
         self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
 
@@ -84,6 +110,17 @@ class MultiHeadAttention(nn.Module):
             self.bias_v = None
 
         self.scale: _Scalar = self.head_dim**-0.5
+        self._reset_parameters()
+
+    def _reset_parameters(self) -> None:
+        if self.in_proj_weight is None:
+            return
+
+        nn.init.kaiming_uniform(self.in_proj_weight)
+        if self.in_proj_bias is not None:
+            fan_in, _ = nn.init._dist._calculate_fan_in_and_fan_out(self.in_proj_weight)
+            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+            nn.init.uniform(self.in_proj_bias, -bound, bound)
 
     def forward(
         self,
@@ -97,9 +134,24 @@ class MultiHeadAttention(nn.Module):
         N, q_len = query.shape[:2]
         k_len, v_len = key.shape[1], value.shape[1]
 
-        q = self.q_proj(query)
-        k = self.k_proj(key)
-        v = self.v_proj(value)
+        if self.in_proj_weight is None:
+            q = self.q_proj(query)
+            k = self.k_proj(key)
+            v = self.v_proj(value)
+        else:
+            if query is key and key is value:
+                qkv = F.linear(query, self.in_proj_weight, self.in_proj_bias)
+                q, k, v = lucid.chunk(qkv, 3, axis=-1)
+            else:
+                w_q, w_k, w_v = lucid.chunk(self.in_proj_weight, 3, axis=0)
+                if self.in_proj_bias is not None:
+                    b_q, b_k, b_v = lucid.chunk(self.in_proj_bias, 3, axis=0)
+                else:
+                    b_q = b_k = b_v = None
+
+                q = F.linear(query, w_q, b_q)
+                k = F.linear(key, w_k, b_k)
+                v = F.linear(value, w_v, b_v)
 
         q = q.reshape(N, self.num_heads, q_len, self.head_dim)
         k = k.reshape(N, self.num_heads, k_len, self.head_dim)

lucid/nn/modules/rnn.py

@@ -284,21 +284,16 @@ class RNNBase(nn.Module):
     ) -> None:
         super().__init__()
         self.is_lstm = False
-        cell_kwargs = {}
         nonlinearity = "tanh"
 
         if mode == "RNN_TANH":
-            cell_cls = RNNCell
-            cell_kwargs: dict[str, object] = {"nonlinearity": nonlinearity}
+            pass
         elif mode == "RNN_RELU":
             nonlinearity = "relu"
-            cell_cls = RNNCell
-            cell_kwargs = {"nonlinearity": nonlinearity}
         elif mode == "LSTM":
-            cell_cls = LSTMCell
             self.is_lstm = True
         elif mode == "GRU":
-            cell_cls = GRUCell
+            pass
         else:
             raise ValueError(
                 f"Invalid mode '{mode}'. Supported modes are 'RNN_TANH', "
@@ -315,18 +310,120 @@ class RNNBase(nn.Module):
         self.batch_first = batch_first
         self.dropout = float(dropout)
 
-        layers: list[nn.Module] = []
         for layer in range(num_layers):
             layer_input_size = input_size if layer == 0 else hidden_size
-            layers.append(
-                cell_cls(
-                    input_size=layer_input_size,
-                    hidden_size=hidden_size,
-                    bias=bias,
-                    **cell_kwargs,
+            sqrt_k = 1.0 / (hidden_size**0.5)
+
+            if mode in ("RNN_TANH", "RNN_RELU"):
+                w_ih = nn.Parameter(
+                    lucid.random.uniform(
+                        -sqrt_k, sqrt_k, (hidden_size, layer_input_size)
+                    )
                 )
-            )
-        self.layers = nn.ModuleList(layers)
+                w_hh = nn.Parameter(
+                    lucid.random.uniform(-sqrt_k, sqrt_k, (hidden_size, hidden_size))
+                )
+            elif mode == "LSTM":
+                w_ih = nn.Parameter(
+                    lucid.random.uniform(
+                        -sqrt_k, sqrt_k, (4 * hidden_size, layer_input_size)
+                    )
+                )
+                w_hh = nn.Parameter(
+                    lucid.random.uniform(
+                        -sqrt_k, sqrt_k, (4 * hidden_size, hidden_size)
+                    )
+                )
+            else:
+                w_ih = nn.Parameter(
+                    lucid.random.uniform(
+                        -sqrt_k, sqrt_k, (3 * hidden_size, layer_input_size)
+                    )
+                )
+                w_hh = nn.Parameter(
+                    lucid.random.uniform(
+                        -sqrt_k, sqrt_k, (3 * hidden_size, hidden_size)
+                    )
+                )
+
+            self.register_parameter(f"weight_ih_l{layer}", w_ih)
+            self.register_parameter(f"weight_hh_l{layer}", w_hh)
+
+            if bias:
+                b_ih = nn.Parameter(
+                    lucid.random.uniform(-sqrt_k, sqrt_k, w_ih.shape[0])
+                )
+                b_hh = nn.Parameter(
+                    lucid.random.uniform(-sqrt_k, sqrt_k, w_hh.shape[0])
+                )
+                self.register_parameter(f"bias_ih_l{layer}", b_ih)
+                self.register_parameter(f"bias_hh_l{layer}", b_hh)
+            else:
+                self.register_parameter(f"bias_ih_l{layer}", None)
+                self.register_parameter(f"bias_hh_l{layer}", None)
+
+    def _rnn_cell(
+        self,
+        input_: Tensor,
+        hx: Tensor,
+        w_ih: Tensor,
+        w_hh: Tensor,
+        b_ih: Tensor | None,
+        b_hh: Tensor | None,
+    ) -> Tensor:
+        hy = F.linear(input_, w_ih, b_ih)
+        hy += F.linear(hx, w_hh, b_hh)
+
+        if self.mode == "RNN_TANH":
+            return F.tanh(hy)
+
+        return F.relu(hy)
+
+    def _lstm_cell(
+        self,
+        input_: Tensor,
+        hx: Tensor,
+        cx: Tensor,
+        w_ih: Tensor,
+        w_hh: Tensor,
+        b_ih: Tensor | None,
+        b_hh: Tensor | None,
+    ) -> tuple[Tensor, Tensor]:
+        gates = F.linear(input_, w_ih, b_ih)
+        gates += F.linear(hx, w_hh, b_hh)
+
+        i_t, f_t, g_t, o_t = lucid.split(gates, 4, axis=1)
+        i_t = F.sigmoid(i_t)
+        f_t = F.sigmoid(f_t)
+        g_t = F.tanh(g_t)
+        o_t = F.sigmoid(o_t)
+
+        c_t = f_t * cx + i_t * g_t
+        h_t = o_t * F.tanh(c_t)
+
+        return h_t, c_t
+
+    def _gru_cell(
+        self,
+        input_: Tensor,
+        hx: Tensor,
+        w_ih: Tensor,
+        w_hh: Tensor,
+        b_ih: Tensor | None,
+        b_hh: Tensor | None,
+    ) -> Tensor:
+        input_gates = F.linear(input_, w_ih, b_ih)
+        hidden_gates = F.linear(hx, w_hh, b_hh)
+
+        i_r, i_z, i_n = lucid.split(input_gates, 3, axis=1)
+        h_r, h_z, h_n = lucid.split(hidden_gates, 3, axis=1)
+
+        r_t = F.sigmoid(i_r + h_r)
+        z_t = F.sigmoid(i_z + h_z)
+        n_t = F.tanh(i_n + r_t * h_n)
+        h_t = (1 - z_t) * n_t + z_t * hx
+
+        return h_t
 
     def _init_hidden(
         self, batch_size: int, dtype: Numeric, device: _DeviceType
@@ -441,7 +538,12 @@ class RNNBase(nn.Module):
         h_n_list: list[Tensor] = []
         c_n_list: list[Tensor] | None = [] if self.is_lstm else None
 
-        for layer_idx, cell in enumerate(self.layers):
+        for layer_idx in range(self.num_layers):
+            w_ih = getattr(self, f"weight_ih_l{layer_idx}")
+            w_hh = getattr(self, f"weight_hh_l{layer_idx}")
+            b_ih = getattr(self, f"bias_ih_l{layer_idx}")
+            b_hh = getattr(self, f"bias_hh_l{layer_idx}")
+
             if self.is_lstm:
                 h_t = hx_h[layer_idx]
                 c_t = hx_c[layer_idx]
@@ -481,9 +583,13 @@ class RNNBase(nn.Module):
                     offset += bs
 
                     if self.is_lstm:
-                        h_t, c_t = cell(step_input, (h_t, c_t))
+                        h_t, c_t = self._lstm_cell(
+                            step_input, h_t, c_t, w_ih, w_hh, b_ih, b_hh
+                        )
+                    elif self.mode == "GRU":
+                        h_t = self._gru_cell(step_input, h_t, w_ih, w_hh, b_ih, b_hh)
                     else:
-                        h_t = cell(step_input, h_t)
+                        h_t = self._rnn_cell(step_input, h_t, w_ih, w_hh, b_ih, b_hh)
 
                     outputs.append(h_t)
                     prev_bs = bs
@@ -512,11 +618,17 @@ class RNNBase(nn.Module):
 
             else:
                 for t in range(seq_len):
+                    step_input = layer_input[t]
                     if self.is_lstm:
-                        h_t, c_t = cell(layer_input[t], (h_t, c_t))
+                        h_t, c_t = self._lstm_cell(
+                            step_input, h_t, c_t, w_ih, w_hh, b_ih, b_hh
+                        )
+                        outputs.append(h_t.unsqueeze(axis=0))
+                    elif self.mode == "GRU":
+                        h_t = self._gru_cell(step_input, h_t, w_ih, w_hh, b_ih, b_hh)
                         outputs.append(h_t.unsqueeze(axis=0))
                     else:
-                        h_t = cell(layer_input[t], h_t)
+                        h_t = self._rnn_cell(step_input, h_t, w_ih, w_hh, b_ih, b_hh)
                         outputs.append(h_t.unsqueeze(axis=0))
 
                 layer_output = lucid.concatenate(tuple(outputs), axis=0)

lucid/nn/modules/sparse.py

@@ -21,7 +21,17 @@ class Embedding(nn.Module):
         super().__init__()
         self.num_embeddings = num_embeddings
         self.embedding_dim = embedding_dim
-        self.padding_idx = padding_idx
+
+        if padding_idx is None:
+            self.padding_idx = None
+        else:
+            pad = int(padding_idx)
+            if pad < 0:
+                pad += num_embeddings
+            if pad < 0 or pad >= num_embeddings:
+                raise IndexError("padding_idx out of range.")
+            self.padding_idx = pad
+
         self.max_norm = max_norm
         self.norm_type = norm_type
 
@@ -32,6 +42,9 @@ class Embedding(nn.Module):
         else:
             self.weight = nn.Parameter(_weight)
 
+        if self.padding_idx is not None:
+            self.weight.data[self.padding_idx] = 0
+
     def forward(self, input_: Tensor) -> Tensor:
         return F.embedding(
             input_, self.weight, self.padding_idx, self.max_norm, self.norm_type
@@ -39,3 +52,5 @@ class Embedding(nn.Module):
 
     def reset_parameters(self) -> None:
         self.weight.data = lucid.random.uniform(-0.1, 0.1, self.weight.shape)
+        if self.padding_idx is not None:
+            self.weight.data[self.padding_idx] = 0

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lucid-dl
-Version: 2.12.0
+Version: 2.12.1
 Summary: Lumerico's Comprehensive Interface for Deep Learning
 Home-page: https://github.com/ChanLumerico/lucid
 Author: ChanLumerico
@@ -48,6 +48,10 @@ Whether you're a student, educator, or an advanced researcher seeking to demysti
 
 ### 🔥 What's New
 
+- New Tensor utility function added: `lucid.Tensor.expand`
+
+- Added Type-Generic Tensors: `lucid.LongTensor`, `lucid.DoubleTensor`, etc.
+
 - Added new visual tool: `lucid.visual.build_tensor_mermaid_chart` which builds a Mermaid chart of given tensor's computatoinal graph
 
 - Added additional `nn.Module` hooks for richer introspection during training:

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

@@ -1,4 +1,4 @@
-lucid/__init__.py,sha256=EwNZkKALNS54RQw_D_1BtqbBo3kFMi4uSveSs6i-bdM,9161
+lucid/__init__.py,sha256=p5SmXNcVcuTeiCmgpKRU6ttNp-Nut58VblUy6B2a5HU,9159
 lucid/error.py,sha256=qnTiVuZm3c5-DIt-OOyobZ7RUm7E1K4NR0j998LG1ug,709
 lucid/port.py,sha256=Kt1YaSWef_eKF4KRj-UFhirvFC5urEESfYQ_BSlBZGE,3811
 lucid/types.py,sha256=Zdz2r4ledouEG-6Gi6yEza5vSLyyTzZJn7AcRKbxy8o,6906
@@ -12,11 +12,11 @@ lucid/_func/ufunc.py,sha256=AnCSykuYC0fNNoZso-TM60Rlq_4c54uMYOG3BHdvy20,30261
 lucid/_fusion/__init__.py,sha256=SVzLiFzs4m1mMOpefKDLFkYqV0zV5FGwFd9hEbUZtSo,68
 lucid/_fusion/base.py,sha256=d6nWuPjYxkie9Xrtbj3JVusnIN61PIoSFFSthJNm9os,3821
 lucid/_fusion/func.py,sha256=9tXzB-QNrx_AvNJiPto807faXKlzjuMG4o9gRgI5usc,1659
-lucid/_tensor/__init__.py,sha256=wFWAMhTnQwThNiBEIT4fcw4ryIm8A4AoR-m9KDhklOQ,40
-lucid/_tensor/base.py,sha256=pEeL4O-R_xFbugdBG8mWipP_pEhtIhWPSgAgHc51ZcQ,5262
-lucid/_tensor/tensor.py,sha256=BLqpVqxP657W8j-oq7la9IVzH5RZnvQDKEMsU1s5rjM,14821
-lucid/_util/__init__.py,sha256=NgOleItHJGVLdJlKHKfpzuSl3vofzJpNsZByHAYJmKs,6838
-lucid/_util/func.py,sha256=ZODVlGdAejMTJnwEls7yMCI5WJ9Thkb3RIq1WwQCS4E,44239
+lucid/_tensor/__init__.py,sha256=prEYQ-GevlGoO7JblW_1dpwarRpCoWzrOYMjWAjSg2Q,180
+lucid/_tensor/base.py,sha256=KWYTFIuc8_ZOeVPUZh-ogT3C9Ey6wA6WLkwepE_CGhk,5323
+lucid/_tensor/tensor.py,sha256=gb9h30Kh6mdKg3eH0XFNLGlhUgyhKF-1lNr_wNcKw1k,19174
+lucid/_util/__init__.py,sha256=_p_qhtPKWzV-kwBpCqSsgXrWSbWWw3I7PMMhgFqiLYg,7085
+lucid/_util/func.py,sha256=DDA6vUYc9b8FeX1pvfI1jW8BpGOm7Wr8E_9btGrYYs0,46921
 lucid/autograd/__init__.py,sha256=hDoK_B2chRFVhoxsT4vxRKangzBEMWqF8gj2hdoTenk,6775
 lucid/data/__init__.py,sha256=qrDIQsnix5ZUEa0yrtomaaWbNJyJ3xEr2gdhRvg70_8,118
 lucid/data/_base.py,sha256=RM8xpBl8qFhm19n7eER_jOsRaxkL3rbOkwUvn6VetSE,5921
@@ -29,7 +29,7 @@ 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/__init__.py,sha256=0BoSSrffJnK3Vqz3yVtvUReuB5uFWTYG8NosK2dz97U,111
 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
@@ -76,6 +76,8 @@ 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/models/seqclf/__init__.py,sha256=qpzGjlHlqe7oQO4KBiz2XtchpoI9u1PUlaPAIh6EY0w,20
+lucid/models/seqclf/bert.py,sha256=wlnZsNci9dMd3yYCY3QoJqgA0s7gnSO2XcfR99l1JaA,990
 lucid/nn/__init__.py,sha256=nyy6px1CxfchWUh68xCiQSxD7Gk65vamhWK8ztRvH68,184
 lucid/nn/fused.py,sha256=75fcXuo6fHSO-JtjuKhowhHSDr4qc5871WR63sUzH0g,5492
 lucid/nn/module.py,sha256=_EWtGkAuWWCPZ5f3t5pJOOzpi14gQBpP7JW2S8o4_GE,26855
@@ -84,7 +86,7 @@ lucid/nn/_kernel/__init__.py,sha256=n1bnYdeb_bNDBKASWGywTRa0Ne9hMAkal3AuVZJgovI,
 lucid/nn/_kernel/activation.py,sha256=mfe48Aw3_Hv0hZEVC7DxDw19XK9XSLfdCOvo2JcZz_o,5662
 lucid/nn/_kernel/attention.py,sha256=1k0gboLObMNVow2v3TwliXC_2v8uKf2o8jHYFuyQqcg,3699
 lucid/nn/_kernel/conv.py,sha256=TiY3EkUAmwFCI1aA8YVMoZJHIRrqmJAXZEPh1C7lons,16412
-lucid/nn/_kernel/embedding.py,sha256=uf_G0aKphxBEOtfR9DHXmSyqkDCwc6tJEBNMKt4CNOU,2390
+lucid/nn/_kernel/embedding.py,sha256=w90-SSr_DYzcI-zLkvye8P2o9C103imPPe4HBRPKUSg,2480
 lucid/nn/_kernel/loss.py,sha256=UD0B5DZ3R98OPZUigHsctL0eAJch2rKQpn1uaI3fzGg,13935
 lucid/nn/_kernel/norm.py,sha256=261WtixerLxFISIroQw8l_zZ3X0b4c_eDy8QHHA-i4M,11992
 lucid/nn/_kernel/pool.py,sha256=IQh5hfKU4PUvnGS1ayorUmytB_vCSxcbAwBYlFKw0iI,10697
@@ -98,12 +100,12 @@ lucid/nn/functional/_loss.py,sha256=b6KT8SrKe5lgAqlAmQnT00Hk7tvd-UcBPNryGYtTPWQ,
 lucid/nn/functional/_norm.py,sha256=yunKJttd3WTxXvzKuugL2LgHLmp-9dMxhHgQ9myLUzA,5041
 lucid/nn/functional/_pool.py,sha256=u6ykwqTZ38b9QPwUqFXpnPhOx2cc_9x9AfH0k26Y9pQ,4085
 lucid/nn/functional/_spatial.py,sha256=lazoSvVMFcauBWRbMOqmkgixA5bDes6scGHVWCgVmHE,3911
-lucid/nn/functional/_util.py,sha256=I2MvzuqPqZBC7Xo_rOq5d1R-f1Hqf7CtTKE06nudY60,5060
+lucid/nn/functional/_util.py,sha256=gfsoGo7JgCHtPkcQqqisO8MfyyK57Pzy7Oeny8k2KKo,5936
 lucid/nn/init/__init__.py,sha256=YFi-HD2TEglweJ-gyX3n4UVZYzd70gcUi1dBu6hnOAY,1533
 lucid/nn/init/_dist.py,sha256=Tj9SKl43ZrJdv99X5qXUowdcts4f4D3tUk7RBmX5uCg,2462
 lucid/nn/modules/__init__.py,sha256=mol5Gfy-3ab5hBYZRxX0vjiI0w5VyKtBxVwj_vrOAZs,285
 lucid/nn/modules/activation.py,sha256=CpiKpzgZHoCp8UO5taCJ9BuwFz5mYUs0o1_TQcEwQbQ,2823
-lucid/nn/modules/attention.py,sha256=pZi7IGsNFu2xCmeLMuyWgveMyi2QXtaKRKQ70yAeE0c,4407
+lucid/nn/modules/attention.py,sha256=XdOrGsS0zTPM8isP7MXoelGuuosNXO9HgD53wGPMBdM,6465
 lucid/nn/modules/conv.py,sha256=KbtInQgKSw3U_qXiqy7x53DZM9YAMUq7sFas1nV7NxY,13932
 lucid/nn/modules/drop.py,sha256=8127XhAbwk0nHWKVcGYqnnzsfmYn-WZ8iR6DXW_al5g,2127
 lucid/nn/modules/einops.py,sha256=3NGbfcBq9PZ9Vlbai53eBGGY4ckeWGXTCdPD73zuuNE,512
@@ -111,8 +113,8 @@ 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=L2rqFRcdr0U33YFeVvthDwDFIE98PrO-OjFiX9IzlIs,21098
-lucid/nn/modules/sparse.py,sha256=EpjiviED2nI55wUjh1twFwa4Lvlrzw0TR6lpCDGeSbo,1147
+lucid/nn/modules/rnn.py,sha256=y_dfvs-2PabKzug9jxMn8o4ir0KfMg64nOS2UKScBZY,24855
+lucid/nn/modules/sparse.py,sha256=C6Kz6Vhe9ko0Ym6JJHOH3HKdNFZJ_xTt2KSRqITXGl8,1620
 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
@@ -136,8 +138,8 @@ 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.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_dl-2.12.1.dist-info/licenses/LICENSE,sha256=vxRFYnVD1IeYtsvw-KmoElfqrjxKHv1h9YTvsG54loQ,1065
+lucid_dl-2.12.1.dist-info/METADATA,sha256=QsDRFeh22Zlxi7RNDDb9aSM75Ly0lZVBtpX82v34kl8,11817
+lucid_dl-2.12.1.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+lucid_dl-2.12.1.dist-info/top_level.txt,sha256=uzP_qBx9iNWIHKJRlElYcBLYVqMpdm9Q1Ma63QPYbFc,6
+lucid_dl-2.12.1.dist-info/RECORD,,