pytensorforge 0.1.0__py3-none-any.whl

src/initializers/RandomUniform.py

@@ -0,0 +1,14 @@
+import numpy as np
+
+class RandomUniform:
+
+    def __init__(self, low=-0.05, high=0.05):
+        self.low = low
+        self.high = high
+
+    def __call__(self, shape):
+        return np.random.uniform(
+            self.low,
+            self.high,
+            shape,
+        ).astype(np.float32)

src/initializers/Zeros.py

@@ -0,0 +1,8 @@
+import numpy as np
+
+from src.initializers.Initializer import Initializer
+
+
+class Zeros(Initializer):
+    def __call__(self, shape):
+        return np.zeros(shape, dtype=np.float32)

src/initializers/__init__.py

@@ -0,0 +1,17 @@
+from .GlorotUniform import GlorotUniform
+from .HeUniform import HeUniform
+from .LecunUniform import LecunUniform
+from .Ones import Ones
+from .Orthogonal import Orthogonal
+from .RandomNormal import RandomNormal
+from .Zeros import Zeros
+
+initializer_fns = {
+    "glorot_uniform": GlorotUniform(),
+    "he_uniform": HeUniform(),
+    "lecun_uniform": LecunUniform(),
+    "zeros": Zeros(),
+    "orthogonal": Orthogonal(),
+    "random_normal": RandomNormal(),
+    "ones": Ones(),
+}

src/loss/CategoricalCrossEntropy.py

@@ -0,0 +1,9 @@
+class CategoricalCrossEntropy:
+
+    def __call__(self, prediction, target):
+
+        eps = 1e-7
+
+        prediction = prediction.clip(eps, 1)
+
+        return -(target * prediction.log()).sum(axis=1).mean()

src/loss/CrossEntropyLoss.py

@@ -0,0 +1,34 @@
+from src.core.Tensor import Tensor
+
+
+class CrossEntropyLoss:
+
+    def __call__(self, logits, target):
+
+        probs = logits.softmax()
+
+        batch = logits.shape[0]
+
+        indices = target.data.astype(int)
+
+        p = probs[
+            Tensor.arange(batch),
+            indices,
+        ]
+
+        return -(p.log()).mean()
+
+# class CrossEntropyLoss:
+#
+#     def __call__(self, logits, target):
+#
+#         probs = logits.softmax()
+#
+#         batch = logits.shape[0]
+#
+#         loss = 0
+#
+#         for i in range(batch):
+#             loss += -probs[i][int(target.data[i])].log()
+#
+#         return loss / batch

src/loss/CrossEntropyWithLogitsLoss.py

@@ -0,0 +1,59 @@
+import numpy as np
+
+from src.core.Tensor import Tensor
+
+
+class CrossEntropyWithLogitsLoss:
+
+    def __call__(self, logits, target):
+
+        batch = logits.shape[0]
+
+        shifted = logits.data - np.max(
+            logits.data,
+            axis=-1,
+            keepdims=True,
+        )
+
+        logsumexp = np.log(
+            np.sum(np.exp(shifted), axis=-1)
+        )
+
+        indices = target.data.astype(np.int64)
+
+        loss = (
+            -shifted[np.arange(batch), indices]
+            + logsumexp
+        ).mean()
+
+        out = Tensor(
+            loss,
+            requires_grad=logits.requires_grad,
+            parents=(logits,),
+            op="CrossEntropyWithLogits",
+        )
+
+        def _backward():
+
+            if not logits.requires_grad:
+                return
+
+            exp = np.exp(shifted)
+
+            probs = exp / np.sum(
+                exp,
+                axis=-1,
+                keepdims=True,
+            )
+
+            grad = probs
+
+            grad[np.arange(batch), indices] -= 1
+
+            grad /= batch
+
+            logits.grad += grad * out.grad
+
+        out._backward = _backward
+
+        return out

src/loss/Hinge.py

@@ -0,0 +1,5 @@
+class Hinge:
+
+    def __call__(self, prediction, target):
+
+        return (1 - target * prediction).maximum(0).mean()

src/loss/Huber.py

@@ -0,0 +1,22 @@
+class Huber:
+
+    def __init__(self, delta=1.0):
+        self.delta = delta
+
+    def __call__(self, prediction, target):
+
+        error = prediction - target
+
+        abs_error = error.abs()
+
+        quadratic = 0.5 * error * error
+
+        linear = self.delta * (
+            abs_error - 0.5 * self.delta
+        )
+
+        return abs_error.where(
+            abs_error < self.delta,
+            quadratic,
+            linear,
+        ).mean()

src/loss/Loss.py

@@ -0,0 +1,6 @@
+class Loss:
+    def forward(self, pred, target):
+        raise NotImplementedError
+
+    def backward(self, pred, target):
+        raise NotImplementedError

src/loss/MSE.py

@@ -0,0 +1,7 @@
+class MSE:
+
+    def __call__(self, prediction, target):
+
+        diff = prediction - target
+
+        return (diff * diff).mean()

src/loss/MSELoss.py

@@ -0,0 +1,10 @@
+import numpy as np
+from src.loss.Loss import Loss
+
+class MSELoss(Loss):
+
+    def forward(self, pred, target):
+        return 0.5 * np.mean((pred - target) ** 2)
+
+    def backward(self, pred, target):
+        return pred - target

src/loss/SparseCategoricalCrossEntropy.py

@@ -0,0 +1,15 @@
+import numpy as np
+
+
+class SparseCategoricalCrossEntropy:
+
+    def __call__(self, prediction, target):
+
+        prediction = prediction.softmax()
+
+        log_probs = prediction.log()
+
+        return -log_probs[
+            np.arange(len(target)),
+            target,
+        ].mean()

src/loss/__init__.py

@@ -0,0 +1,18 @@
+from .CrossEntropyWithLogitsLoss import CrossEntropyWithLogitsLoss
+from .MSE import MSE
+from .CategoricalCrossEntropy import CategoricalCrossEntropy
+from .SparseCategoricalCrossEntropy import SparseCategoricalCrossEntropy
+from .bce import BinaryCrossEntropy
+from .mae import MAE
+from .CrossEntropyLoss import CrossEntropyLoss
+
+losses = {
+    "mse": MSE(),
+    "mean_squared_error": MSE(),
+    "mean_absolute_error": MAE(),
+    "categorical_crossentropy": CategoricalCrossEntropy(),
+    "binary_crossentropy": BinaryCrossEntropy(),
+    "cross_entropy": CrossEntropyLoss(),
+    "sparse_categorical_crossentropy": SparseCategoricalCrossEntropy(),
+    "cross_entropy_with_logits": CrossEntropyWithLogitsLoss()
+}

src/loss/bce.py

@@ -0,0 +1,34 @@
+from src.loss.Loss import Loss
+import numpy as np
+
+class BCELoss(Loss):
+
+    def forward(self, pred, target):
+        eps = 1e-8
+        pred = np.clip(pred, eps, 1 - eps)
+
+        return -np.mean(
+            target * np.log(pred) +
+            (1 - target) * np.log(1 - pred)
+        )
+
+    def backward(self, pred, target):
+        eps = 1e-8
+        pred = np.clip(pred, eps, 1 - eps)
+
+        return (pred - target) / (pred * (1 - pred))
+
+class BinaryCrossEntropy:
+
+    def __call__(self, prediction, target):
+
+        eps = 1e-7
+
+        prediction = prediction.clip(eps, 1 - eps)
+
+        loss = -(
+            target * prediction.log()
+            + (1 - target) * (1 - prediction).log()
+        )
+
+        return loss.mean()

src/loss/mae.py

@@ -0,0 +1,16 @@
+from src.loss.Loss import Loss
+import numpy as np
+
+class MAELoss(Loss):
+
+    def forward(self, pred, target):
+        return np.mean(np.abs(pred - target))
+
+    def backward(self, pred, target):
+        return np.sign(pred - target)
+
+class MAE:
+
+    def __call__(self, prediction, target):
+
+        return (prediction - target).abs().mean()

src/math/clip.py

@@ -0,0 +1,37 @@
+import numpy as np
+
+from src.core.Tensor import Tensor
+
+
+class Clip:
+
+    @staticmethod
+    def forward(x, min_value, max_value):
+
+        clipped = np.clip(x.data, min_value, max_value)
+
+        out = Tensor(
+            clipped,
+            requires_grad=x.requires_grad,
+            parents=(x,),
+            op="Clip",
+        )
+
+        def _backward():
+
+            if not x.requires_grad:
+                return
+
+            mask = (
+                (x.data >= min_value)
+                & (x.data <= max_value)
+            ).astype(np.float32)
+
+            x.grad += out.grad * mask
+
+        out._backward = _backward
+
+        return out
+
+    def __call__(self, x, min_value, max_value):
+        return self.forward(x, min_value, max_value)

src/math/exp.py

@@ -0,0 +1,27 @@
+import numpy as np
+
+from src.core.Tensor import Tensor
+
+
+class Exp:
+
+    @staticmethod
+    def forward(x):
+
+        e = np.exp(x.data)
+
+        out = Tensor(
+            e,
+            requires_grad=x.requires_grad,
+            parents=(x,),
+            op="Exp",
+        )
+
+        def _backward():
+
+            if x.requires_grad:
+                x.grad += out.grad * e
+
+        out._backward = _backward
+
+        return out

src/math/log.py

@@ -0,0 +1,25 @@
+import numpy as np
+
+from src.core.Tensor import Tensor
+
+
+class Log:
+
+    @staticmethod
+    def forward(x):
+
+        out = Tensor(
+            np.log(x.data),
+            requires_grad=x.requires_grad,
+            parents=(x,),
+            op="Log",
+        )
+
+        def _backward():
+
+            if x.requires_grad:
+                x.grad += out.grad / x.data
+
+        out._backward = _backward
+
+        return out

src/math/sigmoid.py

@@ -0,0 +1,5 @@
+import numpy as np
+
+
+def sigmoid(x):
+    return 1 / (1 + np.exp(-x))

src/models/embedding/Embedding.py

@@ -0,0 +1,65 @@
+import numpy as np
+
+from src.neural.Layer import Layer
+from src.core.Tensor import Tensor
+
+
+class Embedding(Layer):
+
+    def __init__(
+        self,
+        vocab_size,
+        embedding_dim,
+    ):
+        super().__init__()
+
+        self.vocab_size = vocab_size
+        self.embedding_dim = embedding_dim
+
+        self.weight = None
+
+    def build(self, input_shape):
+
+        self.weight = self.add_weight(
+            shape=(
+                self.vocab_size,
+                self.embedding_dim,
+            ),
+            initializer="random_normal",
+            name="embedding",
+        )
+
+    def call(self, x):
+
+        indices = x.data.astype(np.int64)
+
+        out = Tensor(
+            self.weight.data[indices],
+            requires_grad=True,
+            parents=(self.weight,),
+            op="Embedding",
+        )
+
+        def _backward():
+
+            if not self.weight.requires_grad:
+                return
+
+            np.add.at(
+                self.weight.grad,
+                indices,
+                out.grad,
+            )
+
+        out._backward = _backward
+
+        return out
+
+    def compute_output_shape(self, input_shape):
+        batch, seq = input_shape
+
+        return (
+            batch,
+            seq,
+            self.embedding_dim,
+        )

src/models/gpt/attention.py

@@ -0,0 +1,158 @@
+import threading
+
+import numpy as np
+
+from src.core.Tensor import Tensor, matmul_policy
+from src.neural.Layer import Layer
+
+_mask_cache = {}
+_mask_lock = threading.Lock()
+
+
+def _future_mask(T):
+    with _mask_lock:
+        mask = _mask_cache.get(T)
+
+        if mask is None:
+            mask = np.triu(np.ones((T, T), dtype=bool), k=1)
+            _mask_cache[T] = mask
+
+        return mask
+
+
+def _identity(x):
+    return x
+
+
+def _to_heads(x, B, T, H, hd):
+    return x.reshape(B, T, H, hd).transpose(0, 2, 1, 3)
+
+
+def _from_heads(x, B, T, d):
+    return np.ascontiguousarray(x.transpose(0, 2, 1, 3)).reshape(B, T, d)
+
+
+def causal_attention(q, k, v, n_heads, rope=None):
+    B, T, d = q.shape
+    H = n_heads
+    hd = d // H
+
+    policy = matmul_policy()
+    rnd = _identity if policy is None else policy
+    scale = 1.0 / float(np.sqrt(hd))
+
+    qh = _to_heads(q.data, B, T, H, hd)
+    kh = _to_heads(k.data, B, T, H, hd)
+    vh = _to_heads(v.data, B, T, H, hd)
+
+    cos = sin = None
+
+    if rope is not None:
+        cos, sin = rope.tables(np.arange(T))
+        qh = rope.rotate(qh, cos, sin)
+        kh = rope.rotate(kh, cos, sin)
+
+    qs = rnd(qh)
+    ks = rnd(kh)
+    vs = rnd(vh)
+
+    scores = rnd(np.matmul(qs, ks.transpose(0, 1, 3, 2)))
+    scores *= scale
+
+    if T > 1:
+        scores[..., _future_mask(T)] = -np.inf
+
+    scores -= scores.max(axis=-1, keepdims=True)
+    np.exp(scores, out=scores)
+    scores /= scores.sum(axis=-1, keepdims=True)
+    probs = scores
+
+    ps = rnd(probs)
+    context = rnd(np.matmul(ps, vs))
+
+    out = Tensor(
+        _from_heads(context, B, T, d),
+        requires_grad=q.requires_grad or k.requires_grad or v.requires_grad,
+        parents=(q, k, v),
+        op="CausalAttention",
+    )
+
+    def _backward():
+        d_out = rnd(_to_heads(out.grad, B, T, H, hd))
+
+        if v.requires_grad:
+            d_v = rnd(np.matmul(ps.transpose(0, 1, 3, 2), d_out))
+            v.grad += _from_heads(d_v, B, T, d)
+
+        if not (q.requires_grad or k.requires_grad):
+            return
+
+        d_p = rnd(np.matmul(d_out, vs.transpose(0, 1, 3, 2)))
+        d_s = probs * (d_p - np.sum(d_p * probs, axis=-1, keepdims=True))
+        d_s *= scale
+        d_s = rnd(d_s)
+
+        if q.requires_grad:
+            d_q = rnd(np.matmul(d_s, ks))
+
+            if rope is not None:
+                d_q = rope.rotate(d_q, cos, sin, inverse=True)
+
+            q.grad += _from_heads(d_q, B, T, d)
+
+        if k.requires_grad:
+            d_k = rnd(np.matmul(d_s.transpose(0, 1, 3, 2), qs))
+
+            if rope is not None:
+                d_k = rope.rotate(d_k, cos, sin, inverse=True)
+
+            k.grad += _from_heads(d_k, B, T, d)
+
+    out._backward = _backward
+
+    return out
+
+
+class CausalSelfAttention(Layer):
+
+    def __init__(self, d_model, num_heads, rope=None):
+        super().__init__()
+
+        if d_model % num_heads != 0:
+            raise ValueError("d_model must be divisible by num_heads")
+
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.head_dim = d_model // num_heads
+        self.rope = rope
+
+        self.Wq = None
+        self.Wk = None
+        self.Wv = None
+        self.Wo = None
+
+    def build(self, input_shape):
+        shape = (self.d_model, self.d_model)
+        self.Wq = self.add_weight(shape=shape, initializer="glorot_uniform", name="Wq")
+        self.Wk = self.add_weight(shape=shape, initializer="glorot_uniform", name="Wk")
+        self.Wv = self.add_weight(shape=shape, initializer="glorot_uniform", name="Wv")
+        self.Wo = self.add_weight(shape=shape, initializer="glorot_uniform", name="Wo")
+        self.built = True
+
+    def call(self, x, mask=None):
+        context = causal_attention(x @ self.Wq, x @ self.Wk, x @ self.Wv, self.num_heads, self.rope)
+        return context @ self.Wo
+
+    def state_dict(self):
+        return {
+            "Wq": self.Wq.data.copy(),
+            "Wk": self.Wk.data.copy(),
+            "Wv": self.Wv.data.copy(),
+            "Wo": self.Wo.data.copy(),
+        }
+
+    def load_state_dict(self, state):
+        self.Wq.data[:] = state["Wq"]
+        self.Wk.data[:] = state["Wk"]
+        self.Wv.data[:] = state["Wv"]
+        self.Wo.data[:] = state["Wo"]