lucid-dl 2.11.0__py3-none-any.whl → 2.11.2__py3-none-any.whl

lucid/nn/_kernel/activation.py

@@ -0,0 +1,188 @@
+import functools
+from types import ModuleType
+
+import numpy as np
+
+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
+
+
+def _norm_axis(axis: int, ndim: int) -> int:
+    return axis if axis >= 0 else axis + ndim
+
+
+class softmax_kernel(Operation):
+    def __init__(self, axis: int = -1) -> None:
+        super().__init__()
+        self.axis = axis
+        self._axis = None
+        self._y = None
+
+    def clear(self) -> None:
+        super().clear()
+        self._axis = None
+        self._y = None
+
+    @func_op(n_in=1, n_ret=1, device="cpu")
+    def cpu(self, a: Tensor) -> _FuncOpReturnType:
+        return self._forward(a, lib_=np, device="cpu")
+
+    @func_op(n_in=1, n_ret=1, device="gpu")
+    def gpu(self, a: Tensor) -> _FuncOpReturnType:
+        return self._forward(a, lib_=mx, device="gpu")
+
+    def _forward(
+        self, a: Tensor, lib_: ModuleType, device: _DeviceType
+    ) -> _FuncOpReturnType:
+        axis = _norm_axis(self.axis, a.ndim)
+        max_val = lib_.max(a.data, axis=axis, keepdims=True)
+        exp_x = lib_.exp(a.data - max_val)
+        sum_exp = lib_.sum(exp_x, axis=axis, keepdims=True)
+        y = exp_x / sum_exp
+
+        self._axis = axis
+        self._y = y
+
+        self.result = Tensor(y, device=device)
+        return self.result, functools.partial(self.__grad__, lib_=lib_)
+
+    def __grad__(self, lib_: ModuleType) -> _GradType:
+        if self.result is None or self.result.grad is None:
+            raise RuntimeError("softmax backward called before forward.")
+        if self._y is None or self._axis is None:
+            raise RuntimeError("softmax cached data missing.")
+
+        dy = self.result.grad
+        y = self._y
+        axis = self._axis
+
+        dot = lib_.sum(dy * y, axis=axis, keepdims=True)
+        dx = y * (dy - dot)
+        return dx
+
+
+class sigmoid_kernel(Operation):
+    def __init__(self) -> None:
+        super().__init__()
+        self._y = None
+
+    def clear(self) -> None:
+        super().clear()
+        self._y = None
+
+    @func_op(n_in=1, n_ret=1, device="cpu")
+    def cpu(self, a: Tensor) -> _FuncOpReturnType:
+        return self._forward(a, lib_=np, device="cpu")
+
+    @func_op(n_in=1, n_ret=1, device="gpu")
+    def gpu(self, a: Tensor) -> _FuncOpReturnType:
+        return self._forward(a, lib_=mx, device="gpu")
+
+    def _forward(
+        self, a: Tensor, lib_: ModuleType, device: _DeviceType
+    ) -> _FuncOpReturnType:
+        y = 1.0 / (1.0 + lib_.exp(-a.data))
+        self._y = y
+        self.result = Tensor(y, device=device)
+        return self.result, functools.partial(self.__grad__)
+
+    def __grad__(self) -> _GradType:
+        if self.result is None or self.result.grad is None or self._y is None:
+            raise RuntimeError("sigmoid backward called before forward.")
+
+        dy = self.result.grad
+        y = self._y
+
+        dx = dy * y * (1 - y)
+        return dx
+
+
+class gelu_kernel(Operation):
+    def __init__(self) -> None:
+        super().__init__()
+        self._x = None
+
+    def clear(self) -> None:
+        super().clear()
+        self._x = None
+
+    @func_op(n_in=1, n_ret=1, device="cpu")
+    def cpu(self, a: Tensor) -> _FuncOpReturnType:
+        return self._forward(a, lib_=np, device="cpu")
+
+    @func_op(n_in=1, n_ret=1, device="gpu")
+    def gpu(self, a: Tensor) -> _FuncOpReturnType:
+        return self._forward(a, lib_=mx, device="gpu")
+
+    def _forward(
+        self, a: Tensor, lib_: ModuleType, device: _DeviceType
+    ) -> _FuncOpReturnType:
+        self._x = a.data
+        c = lib_.sqrt(2.0 / lib_.pi)
+        y = 0.5 * a.data * (1.0 + lib_.tanh(c * (a.data + 0.044715 * (a.data**3))))
+
+        self.result = Tensor(y, device=device)
+        return self.result, functools.partial(self.__grad__, lib_=lib_)
+
+    def __grad__(self, lib_: ModuleType) -> _GradType:
+        if self.result is None or self.result.grad is None or self._x is None:
+            raise RuntimeError("gelu backward called before forward.")
+
+        x = self._x
+        dy = self.result.grad
+        c = lib_.sqrt(2.0 / lib_.pi)
+        t = c * (x + 0.044715 * x**3)
+        dt = c * (1 + 3 * 0.044715 * x**2)
+        sech2 = 1.0 / lib_.cosh(t) ** 2
+
+        dx = 0.5 * (1 + lib_.tanh(t)) + 0.5 * x * sech2 * dt
+        return dy * dx
+
+
+class silu_kernel(Operation):
+    def __init__(self) -> None:
+        super().__init__()
+        self._x = None
+        self._sig = None
+
+    def clear(self) -> None:
+        super().clear()
+        self._x = None
+        self._sig = None
+
+    @func_op(n_in=1, n_ret=1, device="cpu")
+    def cpu(self, a: Tensor) -> _FuncOpReturnType:
+        return self._forward(a, lib_=np, device="cpu")
+
+    @func_op(n_in=1, n_ret=1, device="gpu")
+    def gpu(self, a: Tensor) -> _FuncOpReturnType:
+        return self._forward(a, lib_=mx, device="gpu")
+
+    def _forward(
+        self, a: Tensor, lib_: ModuleType, device: _DeviceType
+    ) -> _FuncOpReturnType:
+        self._x = a.data
+        sig = 1.0 / (1.0 + lib_.exp(-a.data))
+        self._sig = sig
+        y = a.data * sig
+
+        self.result = Tensor(y, device=device)
+        return self.result, functools.partial(self.__grad__)
+
+    def __grad__(self) -> _GradType:
+        if (
+            self.result is None
+            or self.result.grad is None
+            or self._x is None
+            or self._sig is None
+        ):
+            raise RuntimeError("silu backward called before forward.")
+
+        dy = self.result.grad
+        sig = self._sig
+        x = self._x
+
+        dx = dy * (sig + x * sig * (1 - sig))
+        return dx

lucid/nn/_kernel/attention.py

@@ -0,0 +1,125 @@
+import functools
+import math
+from types import ModuleType
+
+import numpy as np
+
+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 _make_causal_mask(lib_: ModuleType, L: int, S: int, dtype: object) -> _TensorData:
+    triu = getattr(lib_, "triu", None)
+    ones = getattr(lib_, "ones", None)
+    if triu is None or ones is None:
+        mask = np.triu(np.ones((L, S), dtype=np.float32), k=1)
+        if lib_ is mx:
+            mask = mx.array(mask)
+    else:
+        mask = triu(ones((L, S), dtype=dtype), k=1)
+    return mask * (-1e12)
+
+
+class scaled_dot_product_attention_kernel(Operation):
+    def __init__(
+        self,
+        attn_mask: Tensor | None = None,
+        is_causal: bool = False,
+        scale: float | None = None,
+    ) -> None:
+        super().__init__()
+        self.attn_mask = attn_mask
+        self.is_causal = bool(is_causal)
+        self.scale = scale
+
+        self._q = None
+        self._k = None
+        self._v = None
+        self._attn = None
+        self._scale = None
+
+    def clear(self) -> None:
+        super().clear()
+        self._q = None
+        self._k = None
+        self._v = None
+        self._attn = None
+        self._scale = None
+
+    @func_op(n_in=3, n_ret=1)
+    def cpu(self, q: Tensor, k: Tensor, v: Tensor) -> _FuncOpReturnType:
+        return self._forward(q, k, v, lib_=np, device="cpu")
+
+    @func_op(n_in=3, n_ret=1, device="gpu")
+    def gpu(self, q: Tensor, k: Tensor, v: Tensor) -> _FuncOpReturnType:
+        return self._forward(q, k, v, lib_=mx, device="gpu")
+
+    def _forward(
+        self, q: Tensor, k: Tensor, v: Tensor, lib_: ModuleType, device: _DeviceType
+    ) -> _FuncOpReturnType:
+        qd = q.data
+        kd = k.data
+        vd = v.data
+
+        scale = self.scale
+        if scale is None:
+            scale = 1.0 / math.sqrt(q.shape[-1])
+
+        kt = lib_.swapaxes(kd, -1, -2)
+        scores = lib_.matmul(qd, kt) * scale
+
+        if self.is_causal:
+            L = q.shape[-2]
+            S = k.shape[-2]
+            scores = scores + _make_causal_mask(lib_, L, S, dtype=scores.dtype)
+
+        if self.attn_mask is not None:
+            scores = scores + self.attn_mask.data
+
+        max_val = lib_.max(scores, axis=-1, keepdims=True)
+        exp_x = lib_.exp(scores - max_val)
+        sum_exp = lib_.sum(exp_x, axis=-1, keepdims=True)
+        attn = exp_x / sum_exp
+
+        out = lib_.matmul(attn, vd)
+
+        self._q = qd
+        self._k = kd
+        self._v = vd
+        self._attn = attn
+        self._scale = scale
+
+        self.result = Tensor(out, device=device)
+        return self.result, functools.partial(self.__grad__, lib_=lib_)
+
+    def __grad__(self, lib_: ModuleType) -> _GradType:
+        if self.result is None or self.result.grad is None:
+            raise RuntimeError("attention backward called before forward.")
+
+        if self._attn is None or self._q is None or self._k is None or self._v is None:
+            raise RuntimeError("attention cached data missing.")
+
+        dy = self.result.grad
+        attn = self._attn
+        qd = self._q
+        kd = self._k
+        vd = self._v
+        scale = self._scale if self._scale is not None else 1.0
+
+        attn_t = lib_.swapaxes(attn, -1, -2)
+        dV = lib_.matmul(attn_t, dy)
+
+        v_t = lib_.swapaxes(vd, -1, -2)
+        dA = lib_.matmul(dy, v_t)
+
+        dot = lib_.sum(dA * attn, axis=-1, keepdims=True)
+        dS = attn * (dA - dot)
+
+        dS = dS * scale
+        dQ = lib_.matmul(dS, kd)
+        dK = lib_.matmul(lib_.swapaxes(dS, -1, -2), qd)
+
+        return dQ, dK, dV

lucid/{_backend → nn/_kernel}/conv.py

@@ -9,7 +9,7 @@ import numpy as np
 from lucid._tensor import Tensor
 from lucid._backend.core import (
     Operation,
-    binary_func_op,
+    func_op,
     _FuncOpReturnType,
     _GradType,
 )
@@ -451,7 +451,7 @@ def _conv_backward_input(
     return grad_input
 
 
-class conv_nd(Operation):
+class conv_nd_kernel(Operation):
     def __init__(
         self,
         stride: int | tuple[int, ...] | list[int],
@@ -481,7 +481,7 @@ class conv_nd(Operation):
 
         return stride, padding, dilation
 
-    @binary_func_op()
+    @func_op(n_in=2, n_ret=1)
     def cpu(self, a: Tensor, b: Tensor) -> _FuncOpReturnType:
         _validate_conv_shapes(a, b, self.groups)
         stride, padding, dilation = self._normalize(b)
@@ -490,7 +490,7 @@ class conv_nd(Operation):
         self.result = Tensor(out)
         return self.result, partial(self.__grad__, a=a, b=b, lib_=np)
 
-    @binary_func_op(device="gpu")
+    @func_op(n_in=2, n_ret=1, device="gpu")
     def gpu(self, a: Tensor, b: Tensor) -> _FuncOpReturnType:
         _validate_conv_shapes(a, b, self.groups)
         stride, padding, dilation = self._normalize(b)
@@ -537,12 +537,3 @@ class conv_nd(Operation):
         macs_per_out = C_in_g * _prod(kernel_size)
         out_elems = N * C_out * _prod(tuple(out_dims))
         return out_elems * macs_per_out
-
-
-def conv_nd_op(
-    stride: int | tuple[int, ...] | list[int],
-    padding: int | tuple[int, ...] | list[int],
-    dilation: int | tuple[int, ...] | list[int],
-    groups: int,
-) -> conv_nd:
-    return conv_nd(stride, padding, dilation, groups)

lucid/nn/_kernel/embedding.py

@@ -0,0 +1,72 @@
+import functools
+from types import ModuleType
+
+import numpy as np
+
+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:
+        super().__init__()
+        self._indices = None
+        self._num_embeddings = None
+
+    def clear(self) -> None:
+        super().clear()
+        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")
+
+    @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")
+
+    def _forward(
+        self, indices: Tensor, weight: Tensor, lib_: ModuleType, device: _DeviceType
+    ) -> _FuncOpReturnType:
+        idx = _as_int_array(indices.data, lib_)
+        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_)
+
+    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.")
+
+        grad_out = self.result.grad
+        idx = self._indices.reshape(-1)
+        grad_flat = grad_out.reshape(idx.shape[0], -1)
+
+        if lib_ is np:
+            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]):
+                grad_w = grad_w.at[idx[i]].add(grad_flat[i])
+
+        return None, grad_w