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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lucid-dl
-Version: 2.11.5
+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.11.5 → 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.11.5 → 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.11.5 → 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.11.5 → 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.11.5 → 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.11.5 → 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.1/lucid/datasets/__init__.py

@@ -0,0 +1,2 @@
+from .mnist import *
+from .cifar import *