catniff 0.2.13 → 0.2.14

package/README.md

@@ -1,6 +1,6 @@
 # Catniff
 
-Catniff is an experimental tensor ops library and autograd engine made to be Torch-like (its name is a play on "catnip" and "differentiation"). This project is heavily under development currently, so keep in mind that APIs can be completely unstable and backwards-incompatible.
+Catniff is a small deep learning framework for Javacript, built to be Torch-like, but more direct on tensors and autograd usage like Tinygrad. This project is under development currently, so keep in mind that APIs can be unstable and backwards-incompatible. On a side-note, the name is a play on "catnip" and "differentiation".
 
 ## Setup
 
@@ -80,14 +80,12 @@ All available APIs are in [`./src/core.ts`](./src/core.ts) if you want to dig de
 
 ## Todos
 
-I'm mostly just learning and playing with this currently, so there are no concrete plans yet, but here is what I currently have in mind:
-
-* Fix whatever is the problem right now (there are a lot of problems right now lol).
-* Add more tensor ops.
-* Proper documentation.
+* Bug fixes.
+* More tensor ops.
+* More detailed documentation.
 * GPU acceleration.
 * Some general neural net APIs.
-* Refactor code.
+* Code refactoring.
 * Proper tests.
 * Option to load more backends.
 

package/dist/core.d.ts

@@ -137,6 +137,7 @@ export declare class Tensor {
     t(): Tensor;
     dot(other: TensorValue | Tensor): Tensor;
     mm(other: TensorValue | Tensor): Tensor;
+    bmm(other: TensorValue | Tensor): Tensor;
     mv(other: TensorValue | Tensor): Tensor;
     matmul(other: TensorValue | Tensor): Tensor;
     static full(shape: number[], num: number, options?: TensorOptions): Tensor;

package/dist/core.js

@@ -1049,6 +1049,64 @@ class Tensor {
         }
         return out;
     }
+    // Batched 3D tensor matmul
+    bmm(other) {
+        other = Tensor.forceTensor(other);
+        // Verify 3D shape
+        if (this.shape.length !== 3 || other.shape.length !== 3 || this.shape[0] !== other.shape[0]) {
+            throw new Error("Inputs are not 3D tensors with the same first dim size");
+        }
+        // Simple matrix multiplication
+        const batchA = this.value;
+        const batchB = other.value;
+        const batchAStrides = this.strides;
+        const batchBStrides = other.strides;
+        const batchSize = this.shape[0];
+        const batchARows = this.shape[1];
+        const batchACols = this.shape[2];
+        const batchBRows = other.shape[1];
+        const batchBCols = other.shape[2];
+        if (batchACols !== batchBRows)
+            throw new Error("Invalid matrices shape for multiplication");
+        const batchCShape = [batchSize, batchARows, batchBCols];
+        const batchCStrides = Tensor.getStrides(batchCShape);
+        const batchCSize = Tensor.shapeToSize(batchCShape);
+        const batchC = new Array(batchCSize).fill(0);
+        for (let q = 0; q < batchSize; q++) {
+            for (let i = 0; i < batchARows; i++) {
+                for (let j = 0; j < batchBCols; j++) {
+                    for (let k = 0; k < batchACols; k++) {
+                        // Tensor values are 1D arrays so we have to get real index using strides
+                        batchC[q * batchCStrides[0] + i * batchCStrides[1] + j * batchCStrides[2]] +=
+                            batchA[q * batchAStrides[0] + i * batchAStrides[1] + k * batchAStrides[2]] *
+                                batchB[q * batchBStrides[0] + k * batchBStrides[1] + j * batchBStrides[2]];
+                    }
+                }
+            }
+        }
+        const out = new Tensor(batchC, { shape: batchCShape, strides: batchCStrides });
+        if (this.requiresGrad) {
+            out.requiresGrad = true;
+            out.children.push(this);
+        }
+        if (other.requiresGrad) {
+            out.requiresGrad = true;
+            out.children.push(other);
+        }
+        if (out.requiresGrad) {
+            out.gradFn = () => {
+                // Disable gradient collecting of gradients themselves
+                const outGrad = out.grad.withGrad(false);
+                const selfNoGrad = this.withGrad(false);
+                const otherNoGrad = other.withGrad(false);
+                if (this.requiresGrad)
+                    Tensor.addGrad(this, outGrad.bmm(otherNoGrad.transpose(1, 2)));
+                if (other.requiresGrad)
+                    Tensor.addGrad(other, selfNoGrad.transpose(1, 2).bmm(outGrad));
+            };
+        }
+        return out;
+    }
     // Convert right-side 1D tensor to a vector (nx1 tensor) to do matmul
     mv(other) {
         other = Tensor.forceTensor(other);

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "catniff",
-  "version": "0.2.13",
-  "description": "A cute autograd engine for Javascript",
+  "version": "0.2.14",
+  "description": "A small Torch-like deep learning framework for Javascript with tensor and autograd support",
   "main": "index.js",
   "scripts": {
     "test": "echo \"Error: no test specified\" && exit 1"
@@ -11,7 +11,6 @@
     "url": "git+https://github.com/nguyenphuminh/catniff.git"
   },
   "keywords": [
-    "cats",
     "catniff",
     "autograd",
     "autodiff",
@@ -27,6 +26,7 @@
     "machine-learning",
     "deep-learning",
     "micrograd",
+    "tinygrad",
     "torch",
     "pytorch"
   ],