lucid-dl 2.12.0__tar.gz → 2.12.1__tar.gz

{lucid_dl-2.12.0/lucid_dl.egg-info → lucid_dl-2.12.1}/PKG-INFO

@@ -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 → lucid_dl-2.12.1}/README.md

@@ -20,6 +20,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 → lucid_dl-2.12.1}/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_dl-2.12.1/lucid/_tensor/__init__.py

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

{lucid_dl-2.12.0 → lucid_dl-2.12.1}/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_dl-2.12.0 → lucid_dl-2.12.1}/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_dl-2.12.0 → lucid_dl-2.12.1}/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_dl-2.12.0 → lucid_dl-2.12.1}/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_dl-2.12.0 → lucid_dl-2.12.1}/lucid/models/__init__.py

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

lucid_dl-2.12.1/lucid/models/seqclf/__init__.py

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

lucid_dl-2.12.1/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_dl-2.12.0 → lucid_dl-2.12.1}/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_dl-2.12.0 → lucid_dl-2.12.1}/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