cogforge-engine 0.2.1__tar.gz → 1.0.0__tar.gz

cogforge_engine-1.0.0/LICENSE

@@ -0,0 +1,201 @@
+Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

{cogforge_engine-0.2.1 → cogforge_engine-1.0.0}/PKG-INFO

@@ -1,12 +1,15 @@
 Metadata-Version: 2.4
 Name: cogforge-engine
-Version: 0.2.1
+Version: 1.0.0
 Summary: A custom autograd engine and Transformer block built from scratch.
 Project-URL: Homepage, https://github.com/avikmjd2/cogforge
 Author-email: Avik Majumder <avikmjd2@gmail.com>
 License-File: LICENSE
 Requires-Python: >=3.8
+Requires-Dist: numexpr>=2.8.0
 Requires-Dist: numpy>=1.20.0
+Provides-Extra: cuda
+Requires-Dist: cupy>=12.0.0; extra == 'cuda'
 Description-Content-Type: text/markdown
 
 # cogforge

{cogforge_engine-0.2.1 → cogforge_engine-1.0.0}/cogforge/app.py

@@ -1,12 +1,55 @@
-import numpy as np
+import numpy
+import backend
+# from backend import np,USE_GPU,NO_GRAD,_NUMEXPR_AVAILABLE,USE_NUMEXPR,_ne,_cp
+
+def to_cpu(a):
+    """Return a numpy array regardless of backend."""
+    return backend._cp.asnumpy(a) if (backend.USE_GPU and backend._cp is not None) else numpy.asarray(a)
+
+def scatter_add(target, indices, values):
+    """target[indices] += values."""
+    if backend.USE_GPU:
+        import cupyx
+        cupyx.scatter_add(target, indices, values)
+    else:
+        numpy.add.at(target, indices, values)
+
+
+def set_numexpr(enabled, threads=None):
+    """User switch: set_numexpr(True, threads=8) or set_numexpr(False)."""
+    # global USE_NUMEXPR
+    if enabled and not backend._NUMEXPR_AVAILABLE:
+        raise RuntimeError("numexpr not installed — run: pip install numexpr")
+    backend.USE_NUMEXPR = bool(enabled)
+    if backend.USE_NUMEXPR and threads:
+        backend._ne.set_num_threads(threads)
+    return backend.USE_NUMEXPR
+
+def _stable_softmax(sd):
+    mx = sd.max(axis=-1, keepdims=True)          # reduction stays numpy
+    if backend.USE_NUMEXPR and not backend.USE_GPU:
+        e = backend._ne.evaluate("exp(sd - mx)")
+        s = e.sum(axis=-1, keepdims=True)
+        return backend._ne.evaluate("e / s")
+    e = backend.np.exp(sd - mx)
+    return e / e.sum(axis=-1, keepdims=True)
+
+def _iadd_scaled(dst, src, coef):                # dst += src * coef, in place
+    if backend.USE_NUMEXPR and not backend.USE_GPU:
+        backend._ne.evaluate("dst + src * coef", out=dst)
+    else:
+        dst += src * coef    
+        
+def needGradientHence(state:bool):
+    backend.NO_GRAD = not state
 
 class Tensor:
     def __init__(self, array,children=(),requires_grad = True,typed="compressed"):
-        self.data = np.asarray(array, dtype=np.float32 if typed == "compressed" else np.float64)
+        self.data = backend.np.asarray(array, dtype=backend.np.float32 if typed == "compressed" else backend.np.float64)
         self.shape = self.data.shape
         self.requires_grad = requires_grad
         # if self.requires_grad:
-        self.grad = np.zeros(self.shape,dtype=np.float32 if typed == "compressed" else np.float64)
+        self.grad = None if backend.NO_GRAD else backend.np.zeros(self.shape,dtype=backend.np.float32 if typed == "compressed" else backend.np.float64)
         self._backwards = lambda:None
         self._children = set(children)
         self.typed = typed
@@ -18,12 +61,36 @@ class Tensor:
         
         def _backward():
             self.grad[key] += out.grad
+            # backend.np.add.at(self.grad, key, out.grad)
+            
+        out._backwards = _backward
+        
+        return out
+    
+    @classmethod
+    def cat(cls,tensors:tuple,axis=-1):
+        """
+        Concatenates a sequence of Tensors along a specified axis.
+        """
+        arrays = tuple(t.data for t in tensors)
+        out = Tensor(backend.np.concatenate(arrays,axis=axis),children=tensors)
+        sizes = [a.shape[axis] for a in arrays[:-1]]
+        split_indices = backend.np.cumsum(sizes)
+        
+        
+        def _backward():
+            split_grads = backend.np.split(out.grad, split_indices, axis=axis)
+            
+            for t,g in zip(tensors,split_grads):
+                if t.requires_grad:
+                    t.grad += g
             
         out._backwards = _backward
         
         return out
         
         
+        
     def __add__(self, other):
         other = other if isinstance(other, Tensor) else Tensor(other)
         out_data = self.data + other.data
@@ -32,13 +99,13 @@ class Tensor:
         def _backward():
             if self.shape != out_data.shape:
                 axis = tuple(range(out.grad.ndim - self.data.ndim))
-                self.grad += np.sum(out.grad, axis=axis).reshape(self.shape)
+                self.grad += backend.np.sum(out.grad, axis=axis).reshape(self.shape)
             else:
                 self.grad += out.grad
 
             if other.shape != out_data.shape:
                 axis = tuple(range(out.grad.ndim - other.data.ndim))
-                other.grad += np.sum(out.grad, axis=axis).reshape(other.shape)
+                other.grad += backend.np.sum(out.grad, axis=axis).reshape(other.shape)
             else:
                 other.grad += out.grad
 
@@ -68,14 +135,14 @@ class Tensor:
             grad_self = other.data * out.grad
             if self.shape != out_data.shape:
                 axis = tuple(range(out.grad.ndim - self.data.ndim))
-                grad_self = np.sum(grad_self, axis=axis).reshape(self.shape)
+                grad_self = backend.np.sum(grad_self, axis=axis).reshape(self.shape)
             self.grad += grad_self
 
             # gradient w.r.t. other
             grad_other = self.data * out.grad
             if other.shape != out_data.shape:
                 axis = tuple(range(out.grad.ndim - other.data.ndim))
-                grad_other = np.sum(grad_other, axis=axis).reshape(other.shape)
+                grad_other = backend.np.sum(grad_other, axis=axis).reshape(other.shape)
             other.grad += grad_other
             
         out._backwards = _backward
@@ -90,12 +157,12 @@ class Tensor:
     #     out = Tensor(out_data, children=(self,other))
         
     #     # def _backward():
-    #     #     self.grad += out.grad @ np.transpose(other.data)
-    #     #     other.grad +=  np.transpose(self.data) @ out.grad
+    #     #     self.grad += out.grad @ backend.np.transpose(other.data)
+    #     #     other.grad +=  backend.np.transpose(self.data) @ out.grad
         
     #     def _backward():
-    #         self.grad += out.grad @ np.swapaxes(other.data, -1, -2)
-    #         other.grad += np.swapaxes(self.data, -1, -2) @ out.grad
+    #         self.grad += out.grad @ backend.np.swapaxes(other.data, -1, -2)
+    #         other.grad += backend.np.swapaxes(self.data, -1, -2) @ out.grad
         
     #     out._backwards = _backward    
     #     return out
@@ -112,15 +179,15 @@ class Tensor:
     def __matmul__(self, other):    
         out = Tensor(self.data @ other.data, children=(self, other))
         def _backward():
-            gs = out.grad @ np.swapaxes(other.data, -1, -2)
-            go = np.swapaxes(self.data, -1, -2) @ out.grad
+            gs = out.grad @ backend.np.swapaxes(other.data, -1, -2)
+            go = backend.np.swapaxes(self.data, -1, -2) @ out.grad
             self.grad  += Tensor.unbroadcast(gs, self.data.shape)
             other.grad += Tensor.unbroadcast(go, other.data.shape)
         out._backwards = _backward
         return out
     
     def relu(self):
-        out_data = np.maximum(0,self.data)
+        out_data = backend.np.maximum(0,self.data)
         
         out = Tensor(out_data, children=(self,))
         
@@ -136,7 +203,7 @@ class Tensor:
         curr_shape = self.data.shape
         
         if not self.data.flags['C_CONTIGUOUS']:
-            arr = np.ascontiguousarray(self.data)
+            arr = backend.np.ascontiguousarray(self.data)
         else:
             arr = self.data
         
@@ -154,7 +221,7 @@ class Tensor:
         
         
     def sigmoid(self):
-        out_data = 1 / (1 + np.exp(-self.data))
+        out_data = 1 / (1 + backend.np.exp(-self.data))
         out = Tensor(out_data, children=(self,))
         def _backward():
             local_derivative = out_data * (1 - out_data)
@@ -165,7 +232,7 @@ class Tensor:
         return out
     
     def tanh(self):
-        out_data = np.tanh(self.data)
+        out_data = backend.np.tanh(self.data)
         out = Tensor(out_data,children=(self,))
         def _backward():
             dervitive = 1-(out_data*out_data)
@@ -175,15 +242,15 @@ class Tensor:
         return out
     
     def softmax(self,axis=-1):
-        maxVal = np.max(self.data,axis=axis,keepdims=True)
-        expVal = np.exp(self.data-maxVal)
+        maxVal = backend.np.max(self.data,axis=axis,keepdims=True)
+        expVal = backend.np.exp(self.data-maxVal)
         
-        out_data = expVal/np.sum(expVal,axis=axis,keepdims=True)
+        out_data = expVal/backend.np.sum(expVal,axis=axis,keepdims=True)
         out=Tensor(out_data,(self,))
         
         def _backward():
             g = out.grad
-            self.grad += out_data * (g - np.sum(g * out_data, axis=axis, keepdims=True))
+            self.grad += out_data * (g - backend.np.sum(g * out_data, axis=axis, keepdims=True))
         
         out._backwards = _backward
         return out
@@ -192,10 +259,10 @@ class Tensor:
     @classmethod
     def cross_entropy_loss(cls,predictions, targets):
        
-        probabilities = np.clip(predictions.data, 1e-15, 1 - 1e-15)
+        probabilities = backend.np.clip(predictions.data, 1e-15, 1 - 1e-15)
         batch_size = predictions.shape[0]
         
-        loss_data = - np.sum(targets * np.log(probabilities)) / batch_size
+        loss_data = - backend.np.sum(targets * backend.np.log(probabilities)) / batch_size
         loss = Tensor(loss_data, children=(predictions,))
         
         def _backward():
@@ -222,7 +289,7 @@ class Tensor:
             
         topoSort(self)
         
-        self.grad = np.ones(self.shape,dtype=np.float32 if self.typed == "compressed" else np.float64)
+        self.grad = backend.np.ones(self.shape,dtype=backend.np.float32 if self.typed == "compressed" else backend.np.float64)
         
         for node in reversed(topo):
             node._backwards()
@@ -245,7 +312,7 @@ class Tensor:
             for child in node._children:
                 stack.append((child,False))
                 
-        self.grad = np.ones(self.shape,dtype=np.float32 if self.typed == "compressed" else np.float64)
+        self.grad = backend.np.ones(self.shape,dtype=backend.np.float32 if self.typed == "compressed" else backend.np.float64)
         
         for node in reversed(topo):
             node._backwards()
@@ -292,13 +359,13 @@ class Tensor:
         targets:     (B, V) one-hot
         mask:        (B,) 1.0 for real positions, 0.0 for <PAD>
         """
-        probs = np.clip(predictions.data, 1e-15, 1 - 1e-15)
+        probs = backend.np.clip(predictions.data, 1e-15, 1 - 1e-15)
         mask = mask.reshape(-1, 1)                          # (B, 1)
         n_real = mask.sum()
         n_real = n_real if n_real > 0 else 1.0
 
         # only real rows contribute to the loss value
-        loss_data = -np.sum(mask * targets * np.log(probs)) / n_real
+        loss_data = -backend.np.sum(mask * targets * backend.np.log(probs)) / n_real
         loss = cls(loss_data, children=(predictions,))
 
         def _backward():
@@ -311,31 +378,45 @@ class Tensor:
         return loss
 
     @classmethod
-    def sparse_softmax_cross_entropy(cls, scores, target_ids):
+    def sparse_softmax_cross_entropy_legacy(cls, scores, target_ids):
         """ scores: (B, T, V) logits ; target_ids: (B, T) int """
         z = scores.data - scores.data.max(axis=-1, keepdims=True)
-        e = np.exp(z)
+        e = backend.np.exp(z)
         p = e / e.sum(axis=-1, keepdims=True)
-        N = int(np.prod(scores.shape[:-1]))
-        flat, idx, rows = p.reshape(N, -1), target_ids.reshape(N), np.arange(N)
-        loss = cls(-np.log(np.clip(flat[rows, idx], 1e-15, 1.0)).sum() / N, children=(scores,))
+        N = int(backend.np.prod(scores.shape[:-1]))
+        flat, idx, rows = p.reshape(N, -1), target_ids.reshape(N), backend.np.arange(N)
+        loss = cls(-backend.np.log(backend.np.clip(flat[rows, idx], 1e-15, 1.0)).sum() / N, children=(scores,))
         def _backward():
-            g = p.reshape(N, -1).copy()
+            # g = p.reshape(N, -1).copy()
+            g = p.reshape(N, -1)
             g[rows, idx] -= 1.0                      # (softmax - onehot)
             scores.grad += (g.reshape(scores.shape) / N) * loss.grad
         loss._backwards = _backward
         return loss
     
+    @classmethod
+    def sparse_softmax_cross_entropy(cls, scores, target_ids):
+        p = _stable_softmax(scores.data)
+        N = int(backend.np.prod(scores.shape[:-1]))
+        flat, idx, rows = p.reshape(N, -1), target_ids.reshape(N), backend.np.arange(N)
+        loss = cls(-backend.np.log(backend.np.clip(flat[rows, idx], 1e-15, 1.0)).sum() / N, children=(scores,))
+        def _backward():
+            g = p.reshape(N, -1)
+            g[rows, idx] -= 1.0
+            _iadd_scaled(scores.grad, g.reshape(scores.shape), backend.np.float32(loss.grad / N))
+        loss._backwards = _backward
+        return loss
+    
     #Legacy
     @classmethod
     def softmax_cross_entropy_old(cls, scores, targets):
         z = scores.data
         z = z - z.max(axis=1, keepdims=True)          # stability
-        e = np.exp(z)
+        e = backend.np.exp(z)
         p = e / e.sum(axis=1, keepdims=True)          # softmax, computed privately
         B = scores.shape[0]
 
-        loss = cls(-np.sum(targets * np.log(np.clip(p, 1e-15, 1.0))) / B, children=(scores,))
+        loss = cls(-backend.np.sum(targets * backend.np.log(backend.np.clip(p, 1e-15, 1.0))) / B, children=(scores,))
 
         def _backward():
             scores.grad += (p - targets) / B * loss.grad
@@ -346,10 +427,10 @@ class Tensor:
     def softmax_cross_entropy(cls, scores, targets):
         z = scores.data
         z = z - z.max(axis=-1, keepdims=True)          
-        e = np.exp(z)
+        e = backend.np.exp(z)
         p = e / e.sum(axis=-1, keepdims=True)
-        N = np.prod(scores.shape[:-1])                  # B for 2D, B*T for 3D
-        loss = cls(-np.sum(targets * np.log(np.clip(p, 1e-15, 1.0))) / N, children=(scores,))
+        N = backend.np.prod(scores.shape[:-1])                  # B for 2D, B*T for 3D
+        loss = cls(-backend.np.sum(targets * backend.np.log(backend.np.clip(p, 1e-15, 1.0))) / N, children=(scores,))
         def _backward():
             scores.grad += (p - targets) / N * loss.grad
         loss._backwards = _backward
@@ -360,14 +441,14 @@ class Tensor:
     def softmax_cross_entropy_masked(cls, scores, targets, mask):
         z = scores.data
         z = z - z.max(axis=1, keepdims=True)
-        e = np.exp(z)
+        e = backend.np.exp(z)
         p = e / e.sum(axis=1, keepdims=True)
 
         mask = mask.reshape(-1, 1)                     
         n_real = mask.sum()
         n_real = n_real if n_real > 0 else 1.0
 
-        loss = cls(-np.sum(mask * targets * np.log(np.clip(p, 1e-15, 1.0))) / n_real, children=(scores,))
+        loss = cls(-backend.np.sum(mask * targets * backend.np.log(backend.np.clip(p, 1e-15, 1.0))) / n_real, children=(scores,))
 
         def _backward():
             grad = (p - targets) / n_real
@@ -377,19 +458,19 @@ class Tensor:
         return loss
     
     def transpose(self,axes):
-        out = Tensor(np.transpose(self.data,axes=axes),children=(self,))
-        inv = np.argsort(axes)
+        out = Tensor(backend.np.transpose(self.data,axes=axes),children=(self,))
+        inv = backend.np.argsort(axes)
         def _backward():
-            self.grad+=np.transpose(out.grad,inv)
+            self.grad+=backend.np.transpose(out.grad,inv)
             
         out._backwards = _backward
         return out
     
     def masked_fill(self, mask, value):
-        out_data = np.where(mask, value, self.data)
+        out_data = backend.np.where(mask, value, self.data)
         out = Tensor(out_data, children=(self,))
         def _backward():
-            self.grad += np.where(mask, 0.0, 1.0) * out.grad
+            self.grad += backend.np.where(mask, 0.0, 1.0) * out.grad
         out._backwards = _backward
         return out
                             
@@ -398,10 +479,10 @@ class Tensor:
             
 class Linear:
     def __init__(self, nin, nout):
-        self.W = Tensor(np.random.randn(nin, nout) * np.sqrt(2.0/nin))
-        self.v_W = np.zeros(self.W.shape)
-        self.b = Tensor(np.zeros((nout,)))
-        self.v_b = np.zeros(self.b.shape)
+        self.W = Tensor(backend.np.random.randn(nin, nout) * backend.np.sqrt(2.0/nin))
+        self.v_W = backend.np.zeros(self.W.shape)
+        self.b = Tensor(backend.np.zeros((nout,)))
+        self.v_b = backend.np.zeros(self.b.shape)
         
     def __call__ (self,x):
         if(x.shape[-1]!=self.W.shape[0]):
@@ -436,10 +517,10 @@ class MLP:
             weights_dict[f'W_{i}'] = layer.W.data
             weights_dict[f'b_{i}'] = layer.b.data
             
-        np.savez_compressed(filename, **weights_dict)
+        backend.np.savez_compressed(filename, **weights_dict)
         
     def load(self, filename="best_model.npz"):
-        with np.load(filename) as data:
+        with backend.np.load(filename) as data:
             for i, layer in enumerate(self.layers):
                 layer.W.data = data[f'W_{i}']
                 layer.b.data = data[f'b_{i}']
@@ -448,15 +529,17 @@ class MLP:
         
 class Embedding:
     def __init__(self, vocab_size, embedding_dim):
-        self.weights = Tensor(np.random.randn(vocab_size, embedding_dim))
+        self.weights = Tensor(backend.np.random.randn(vocab_size, embedding_dim))
         
     def __call__(self, input_indices):
-        out_data = self.weights.data[input_indices]
+        idx = backend.np.asarray(input_indices)   
+        out_data = self.weights.data[idx]
         
         out = Tensor(out_data, children=(self.weights,))
         
         def _backward():
-            np.add.at(self.weights.grad,input_indices,out.grad)
+            # backend.np.add.at(self.weights.grad,input_indices,out.grad)
+            scatter_add(self.weights.grad, idx, out.grad)
         
         out._backwards = _backward
         return out
@@ -500,7 +583,7 @@ class RNN:
             h = prev_hidden
         else:
             B = xs[0].shape[0]
-            h = Tensor(np.zeros((B, self.hidden_dim)))
+            h = Tensor(backend.np.zeros((B, self.hidden_dim)))
         
         
         hidden_states = []
@@ -556,7 +639,7 @@ class SGD:
         self.lr = learning_rate
         
     def clip_grads(self, max_norm=5.0):
-        total = np.sqrt(sum(np.sum(p.grad ** 2) for p in self.parameters))
+        total = backend.np.sqrt(sum(backend.np.sum(p.grad ** 2) for p in self.parameters))
         if total > max_norm:
             scale = max_norm / (total + 1e-6)
             for p in self.parameters:
@@ -568,7 +651,7 @@ class SGD:
             
     def zero_grad(self):
         for p in self.parameters:
-            p.grad = np.zeros_like(p.grad)
+            p.grad = backend.np.zeros_like(p.grad)
             
 class Sequential:
     def __init__(self,layers):
@@ -606,11 +689,11 @@ class BatchNorm1D:
         self.momentum = momentum
         self.training = True
         
-        self.gamma = Tensor(np.ones(dim))
-        self.beta = Tensor(np.zeros(dim))
+        self.gamma = Tensor(backend.np.ones(dim))
+        self.beta = Tensor(backend.np.zeros(dim))
         
-        self.running_mean = np.zeros(dim)
-        self.running_var = np.ones(dim)
+        self.running_mean = backend.np.zeros(dim)
+        self.running_var = backend.np.ones(dim)
         
     
     def __call__(self,x:Tensor):
@@ -621,17 +704,17 @@ class BatchNorm1D:
             var = x.data.var(axis=reduce_dims, keepdims=True)
             
             self.x_centered = x.data - mean
-            self.std_inv = 1.0 / np.sqrt(var + self.eps)
+            self.std_inv = 1.0 / backend.np.sqrt(var + self.eps)
             self.x_hat = self.x_centered * self.std_inv
             
-            N = np.prod([x.data.shape[d] for d in reduce_dims])
+            N = backend.np.prod([x.data.shape[d] for d in reduce_dims])
             unbiased_var = var.squeeze() * (N / (N - 1)) if N > 1 else var.squeeze()
             
             self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean.squeeze()
             self.running_var = (1 - self.momentum) * self.running_var + self.momentum * unbiased_var
             
         else:
-            self.x_hat = (x.data-self.running_mean)/np.sqrt(self.running_var + self.eps)
+            self.x_hat = (x.data-self.running_mean)/backend.np.sqrt(self.running_var + self.eps)
         
         out_data = self.gamma.data * self.x_hat + self.beta.data
         
@@ -642,16 +725,16 @@ class BatchNorm1D:
                 return
             
             dout = out.grad
-            self.gamma.grad += np.sum(dout * self.x_hat, axis=reduce_dims)
-            self.beta.grad += np.sum(dout, axis=reduce_dims)
+            self.gamma.grad += backend.np.sum(dout * self.x_hat, axis=reduce_dims)
+            self.beta.grad += backend.np.sum(dout, axis=reduce_dims)
             
-            N = np.prod([x.data.shape[d] for d in reduce_dims])
+            N = backend.np.prod([x.data.shape[d] for d in reduce_dims])
 
             dx_hat = dout * self.gamma.data
             dx = (1.0 / N) * self.std_inv * (
                 N * dx_hat 
-                - np.sum(dx_hat, axis=reduce_dims, keepdims=True) 
-                - self.x_hat * np.sum(dx_hat * self.x_hat, axis=reduce_dims, keepdims=True)
+                - backend.np.sum(dx_hat, axis=reduce_dims, keepdims=True) 
+                - self.x_hat * backend.np.sum(dx_hat * self.x_hat, axis=reduce_dims, keepdims=True)
             )
 
             x.grad += dx
@@ -724,8 +807,8 @@ class Adam:
     def __init__(self, parameters, lr=1e-3, beta1=0.9, beta2=0.999, eps=1e-8):
         self.parameters = parameters
         self.lr = lr; self.beta1 = beta1; self.beta2 = beta2; self.eps = eps
-        self.m = [np.zeros_like(p.data) for p in parameters]
-        self.v = [np.zeros_like(p.data) for p in parameters]
+        self.m = [backend.np.zeros_like(p.data) for p in parameters]
+        self.v = [backend.np.zeros_like(p.data) for p in parameters]
         self.t = 0
 
     def step(self):
@@ -735,14 +818,14 @@ class Adam:
             self.v[i] = self.beta2*self.v[i] + (1-self.beta2)*(p.grad**2)
             m_hat = self.m[i] / (1 - self.beta1**self.t)
             v_hat = self.v[i] / (1 - self.beta2**self.t)
-            p.data -= self.lr * m_hat / (np.sqrt(v_hat) + self.eps)
+            p.data -= self.lr * m_hat / (backend.np.sqrt(v_hat) + self.eps)
 
     def zero_grad(self):
         for p in self.parameters:
-            p.grad = np.zeros_like(p.grad)
+            p.grad = backend.np.zeros_like(p.grad)
 
     def clip_grads(self, max_norm=5.0):
-        total = np.sqrt(sum(np.sum(p.grad**2) for p in self.parameters))
+        total = backend.np.sqrt(sum(backend.np.sum(p.grad**2) for p in self.parameters))
         if total > max_norm:
             scale = max_norm / (total + 1e-6)
             for p in self.parameters:
@@ -751,7 +834,7 @@ class Adam:
                 
 class Attention:
     def __init__(self,dk):
-        self.scale = 1.0/np.sqrt(dk)
+        self.scale = 1.0/backend.np.sqrt(dk)
     
     def __call__(self, Q:Tensor,K:Tensor,V:Tensor,mask=None):
         axes = list(range(Q.data.ndim)); axes[-1], axes[-2] = axes[-2], axes[-1]
@@ -765,12 +848,13 @@ class Attention:
         return intermediate @ V
         
 class MultiHeadAttention:
-    def __init__(self,dinp,dmodel, dout,n):
+    def __init__(self,dinp,dmodel, dout,n,rope=None):
         assert dmodel%n==0
         self.h = n
         self.dk = dmodel//n
         self.q = Linear(dinp,dmodel) #dmodel = dk*h, we will slice it 
         self.k = Linear(dinp,dmodel)
+        self.rope = rope
         self.v = Linear(dinp,dmodel)
         self.o = Linear(dmodel,dout)
         self.attention = Attention(self.dk)
@@ -787,6 +871,9 @@ class MultiHeadAttention:
         Q = self.split(self.q(query))
         K = self.split(self.k(key))
         V = self.split(self.v(value))
+        if self.rope is not None:
+            Q = self.rope(Q)
+            K = self.rope(K)
         out = self.attention(Q=Q,K=K,V=V,mask=mask)
         return self.o(self.merge(out))
     
@@ -798,13 +885,13 @@ class MultiHeadAttention:
 class LayerNorm:
     def __init__(self, dim, eps=1e-5):
         self.eps = eps
-        self.gamma = Tensor(np.ones(dim))      
-        self.beta  = Tensor(np.zeros(dim))     
+        self.gamma = Tensor(backend.np.ones(dim))      
+        self.beta  = Tensor(backend.np.zeros(dim))     
 
     def __call__(self, x):
         mu  = x.data.mean(axis=-1, keepdims=True)
         var = x.data.var(axis=-1, keepdims=True)
-        std_inv = 1.0 / np.sqrt(var + self.eps)
+        std_inv = 1.0 / backend.np.sqrt(var + self.eps)
         x_hat = (x.data - mu) * std_inv
         out_data = self.gamma.data * x_hat + self.beta.data
         out = Tensor(out_data, children=(x, self.gamma, self.beta))
@@ -813,13 +900,13 @@ class LayerNorm:
         def _backward():
             dout = out.grad
             axes = tuple(range(dout.ndim - 1))            
-            self.gamma.grad += np.sum(dout * x_hat, axis=axes)
-            self.beta.grad  += np.sum(dout, axis=axes)
+            self.gamma.grad += backend.np.sum(dout * x_hat, axis=axes)
+            self.beta.grad  += backend.np.sum(dout, axis=axes)
             dxhat = dout * self.gamma.data
             dx = std_inv / D * (
                 D * dxhat
-                - np.sum(dxhat, axis=-1, keepdims=True)
-                - x_hat * np.sum(dxhat * x_hat, axis=-1, keepdims=True)
+                - backend.np.sum(dxhat, axis=-1, keepdims=True)
+                - x_hat * backend.np.sum(dxhat * x_hat, axis=-1, keepdims=True)
             )
             x.grad += dx
         out._backwards = _backward
@@ -842,9 +929,9 @@ class FeedForward:
 class Transformer:
     """The transformer block has no bridge. So demb = dmodel"""
     
-    def __init__(self, dmodel, n, dff = None):
+    def __init__(self, dmodel, n, dff = None,rope=None):
         self.ln1 = LayerNorm(dmodel)
-        self.attn = MultiHeadAttention(dinp=dmodel,dmodel=dmodel,dout=dmodel,n=n)
+        self.attn = MultiHeadAttention(dinp=dmodel,dmodel=dmodel,dout=dmodel,n=n,rope=rope)
         self.ln2 = LayerNorm(dmodel)
         self.ff  = FeedForward(dmodel=dmodel,dff=dff)
         
@@ -863,11 +950,11 @@ class Transformer:
         
 class PositionalEncoding:
     def __init__(self,max_len,dmodel):
-        pe = np.zeros((max_len,dmodel))
-        pos = np.arange(max_len).reshape(-1,1)
-        div = np.exp(np.arange(0, dmodel, 2) * (-np.log(10000.0) / dmodel))
-        pe[:,0::2] = np.sin(pos*div)
-        pe[:,1::2] = np.cos(pos*div)
+        pe = backend.np.zeros((max_len,dmodel))
+        pos = backend.np.arange(max_len).reshape(-1,1)
+        div = backend.np.exp(backend.np.arange(0, dmodel, 2) * (-backend.np.log(10000.0) / dmodel))
+        pe[:,0::2] = backend.np.sin(pos*div)
+        pe[:,1::2] = backend.np.cos(pos*div)
         self.pe = pe
         
     
@@ -877,4 +964,34 @@ class PositionalEncoding:
     
     def parameters(self):
         return []
- 
+        
+class RotatoryPositionalEncoding:
+    def __init__(self,max_len, dim, base=10000.0):
+        assert dim % 2 == 0, "RoPE dim must be even"
+        theta = 1.0/(base**(backend.np.arange(0,dim,2)/float(dim)))
+        m = backend.np.arange(max_len)
+        
+        freq = backend.np.outer(m,theta)
+        cos = backend.np.cos(freq)
+        sin = backend.np.sin(freq)
+        
+        self.cos = backend.np.concatenate([cos, cos], axis=-1)               
+        self.sin = backend.np.concatenate([sin, sin], axis=-1)
+        
+    def xbar(self,x):
+        d = x.shape[-1]
+        half = d // 2
+        
+        x1 = x[...,:half]
+        x2 = x[...,half:]
+        
+        return Tensor.cat((-x2,x1),axis=-1)
+    
+    def __call__(self,x:Tensor):
+        T = x.shape[-2]
+        cos = self.cos[:T]
+        sin = self.sin[:T]
+        return x*cos + self.xbar(x)*sin
+    
+    def parameters(self):
+        return []

cogforge_engine-1.0.0/cogforge/backend.py

@@ -0,0 +1,27 @@
+import numpy 
+try:
+    import numexpr as _ne
+    _NUMEXPR_AVAILABLE = True
+except ImportError:
+    _NUMEXPR_AVAILABLE = False
+    
+USE_NUMEXPR = False  
+
+try:
+    import cupy as _cp
+    _CUPY_AVAILABLE = True
+except ImportError:
+    _cp = None
+    _CUPY_AVAILABLE = False        
+    
+np = numpy
+USE_GPU = False
+NO_GRAD = False
+
+def use_gpu(enabled = True):
+    global np, USE_GPU
+    if enabled and not _CUPY_AVAILABLE:
+        raise RuntimeError("cupy not installed")
+    np = _cp if enabled else numpy
+    USE_GPU = bool(enabled)
+    return USE_GPU

cogforge_engine-1.0.0/cogforge/models.py

@@ -0,0 +1,88 @@
+import backend
+from app import Embedding, PositionalEncoding,Transformer,LayerNorm,Linear,RotatoryPositionalEncoding
+
+class GPTV1:
+    def __init__(self, vocab, d_model, n_heads, n_layers, max_len, d_ff=None):
+        self.tok = Embedding(vocab, d_model)
+        self.pos    = PositionalEncoding(max_len, d_model)
+        self.blocks = [Transformer(dmodel=d_model,n=n_heads,dff=d_ff) for _ in range(n_layers)]
+        self.ln_f   = LayerNorm(d_model) 
+        self.head   = Linear(d_model, vocab)
+        self.max_len = max_len
+        
+    def __call__(self, idx):
+        T = idx.shape[1]
+        casual = backend.np.triu((backend.np.ones((T,T))),k=1).astype(bool)
+        x = self.pos(self.tok(idx))
+        for block in self.blocks:
+            x = block(x,mask=casual)
+            
+        x = self.ln_f(x)
+        return self.head(x)
+    
+    def parameters(self):
+        ps = self.tok.parameters() + self.ln_f.parameters() + self.head.parameters()
+        for b in self.blocks: ps += b.parameters()
+        return ps
+    
+    def generate(self, idx, n_new, temperature=1.0, top_k=None):
+        """
+        idx: (B, T) int array — the prompt (B is usually 1)
+        returns: (B, T + n_new) int array
+        """
+        for _ in range(n_new):
+            cond = idx[:, -self.max_len:]                    # crop to context window
+            logits = self(cond).data[:, -1, :] / temperature # (B, V) — last position only
+            if top_k is not None:
+                kth = backend.np.sort(logits, axis=-1)[:, -top_k][:, None]
+                logits = backend.np.where(logits < kth, -1e9, logits)  # keep only top-k choices
+            z = logits - logits.max(-1, keepdims=True)
+            p = backend.np.exp(z).astype(backend.np.float64); p /= p.sum(-1, keepdims=True)     # softmax → probabilities
+            nxt = backend.np.array([[backend.np.random.choice(len(pr), p=pr)] for pr in p])  # sample
+            idx = backend.np.concatenate([idx, nxt], axis=1)         # append, feed back in
+        return idx                                            # <-- was a bare `return` 
+        
+        
+        
+        
+class GPT2:
+    def __init__(self, vocab, d_model, n_heads, n_layers, max_len, d_ff=None, base=10000.0):
+        assert d_model % n_heads == 0
+        self.tok    = Embedding(vocab, d_model)
+        self.rope   = RotatoryPositionalEncoding(max_len, d_model // n_heads, base)  # dim = d_k
+        self.blocks = [Transformer(d_model, n_heads, dff=d_ff, rope=self.rope)
+                       for _ in range(n_layers)]
+        self.ln_f   = LayerNorm(d_model)
+        self.head   = Linear(d_model, vocab)
+        self.max_len = max_len
+
+
+    def __call__(self, idx):
+        T = idx.shape[1]
+        causal = backend.np.triu(backend.np.ones((T, T)), k=1).astype(bool)
+        x = self.tok(idx)                          # straight from token embedding
+        for block in self.blocks:
+            x = block(x, mask=causal)
+        return self.head(self.ln_f(x))
+
+    def parameters(self):
+        ps = self.tok.parameters() + self.ln_f.parameters() + self.head.parameters()
+        for b in self.blocks: ps += b.parameters()
+        return ps   
+    
+    def generate(self, idx, n_new, temperature=1.0, top_k=None):
+        """
+        idx: (B, T) int array — the prompt (B is usually 1)
+        returns: (B, T + n_new) int array
+        """
+        for _ in range(n_new):
+            cond = idx[:, -self.max_len:]                    # crop to context window
+            logits = self(cond).data[:, -1, :] / temperature # (B, V) — last position only
+            if top_k is not None:
+                kth = backend.np.sort(logits, axis=-1)[:, -top_k][:, None]
+                logits = backend.np.where(logits < kth, -1e9, logits)  # keep only top-k choices
+            z = logits - logits.max(-1, keepdims=True)
+            p = backend.np.exp(z).astype(backend.np.float64); p /= p.sum(-1, keepdims=True)     # softmax → probabilities
+            nxt = backend.np.array([[backend.np.random.choice(len(pr), p=pr)] for pr in p])  # sample
+            idx = backend.np.concatenate([idx, nxt], axis=1)         # append, feed back in
+        return idx                               

{cogforge_engine-0.2.1 → cogforge_engine-1.0.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "cogforge-engine" 
-version = "0.2.1"
+version = "1.0.0"
 authors = [
   { name="Avik Majumder", email="avikmjd2@gmail.com" },
 ]
@@ -13,6 +13,12 @@ readme = "README.md"
 requires-python = ">=3.8"
 dependencies = [
     "numpy>=1.20.0",
+    "numexpr>=2.8.0",
+]
+
+[project.optional-dependencies]
+cuda = [
+    "cupy>=12.0.0", # Optional dependency
 ]
 
 [project.urls]

cogforge_engine-0.2.1/cogforge/models.py

@@ -1,44 +0,0 @@
-import numpy as np
-from app import Embedding, PositionalEncoding,Transformer,LayerNorm,Linear
-
-class GPTV1:
-    def __init__(self, vocab, d_model, n_heads, n_layers, max_len, d_ff=None):
-        self.tok = Embedding(vocab, d_model)
-        self.pos    = PositionalEncoding(max_len, d_model)
-        self.blocks = [Transformer(dmodel=d_model,n=n_heads,dff=d_ff) for _ in range(n_layers)]
-        self.ln_f   = LayerNorm(d_model) 
-        self.head   = Linear(d_model, vocab)
-        self.max_len = max_len
-        
-    def __call__(self, idx):
-        T = idx.shape[1]
-        casual = np.triu((np.ones((T,T))),k=1).astype(bool)
-        x = self.pos(self.tok(idx))
-        for block in self.blocks:
-            x = block(x,mask=casual)
-            
-        x = self.ln_f(x)
-        return self.head(x)
-    
-    def parameters(self):
-        ps = self.tok.parameters() + self.ln_f.parameters() + self.head.parameters()
-        for b in self.blocks: ps += b.parameters()
-        return ps
-    
-    def generate(self, idx, n_new, temperature=1.0, top_k=None):
-        """
-        idx: (B, T) int array — the prompt (B is usually 1)
-        returns: (B, T + n_new) int array
-        """
-        for _ in range(n_new):
-            cond = idx[:, -self.max_len:]                    # crop to context window
-            logits = self(cond).data[:, -1, :] / temperature # (B, V) — last position only
-            if top_k is not None:
-                kth = np.sort(logits, axis=-1)[:, -top_k][:, None]
-                logits = np.where(logits < kth, -1e9, logits)  # keep only top-k choices
-            z = logits - logits.max(-1, keepdims=True)
-            p = np.exp(z).astype(np.float64); p /= p.sum(-1, keepdims=True)     # softmax → probabilities
-            nxt = np.array([[np.random.choice(len(pr), p=pr)] for pr in p])  # sample
-            idx = np.concatenate([idx, nxt], axis=1)         # append, feed back in
-        return idx                                            # <-- was a bare `return` 
-