catniff 0.2.13 → 0.2.15

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);
@@ -1086,19 +1144,94 @@ class Tensor {
     // General matrix multiplication with different shapes
     matmul(other) {
         other = Tensor.forceTensor(other);
-        if (this.shape.length === 1 && other.shape.length === 1) {
+        const isThis1D = this.shape.length === 1;
+        const isOther1D = other.shape.length === 1;
+        if (isThis1D && isOther1D) {
             return this.dot(other);
         }
-        else if (this.shape.length === 1 && other.shape.length === 2) {
+        else if (isThis1D && other.shape.length === 2) {
             return this.unsqueeze(0).mm(other).squeeze(0);
         }
-        else if (this.shape.length === 2 && other.shape.length === 1) {
+        else if (this.shape.length === 2 && isOther1D) {
             return this.mv(other);
         }
         else if (this.shape.length === 2 && other.shape.length === 2) {
             return this.mm(other);
         }
-        // Too lazy for batched matmul
+        else if ((isThis1D && other.shape.length > 2) ||
+            (isOther1D && this.shape.length > 2) ||
+            (other.shape.length > 2 && this.shape.length > 2)) {
+            // Append/prepend dims if needed
+            const self = isThis1D ? this.unsqueeze(0) : this;
+            other = isOther1D ? other.unsqueeze(1) : other;
+            // Padding
+            const [selfStrides, otherStrides, selfShape, otherShape] = Tensor.padShape(self.strides, other.strides, self.shape, other.shape);
+            const lastDim = selfShape.length - 1;
+            // Prepare data for broadcasting
+            const batchA = self.value;
+            const batchB = other.value;
+            const batchARows = selfShape[lastDim - 1];
+            const batchACols = selfShape[lastDim];
+            const batchBRows = otherShape[lastDim - 1];
+            const batchBCols = otherShape[lastDim];
+            // Verify if can do matmul
+            if (batchACols !== batchBRows)
+                throw new Error("Invalid matrices shape for multiplication");
+            // Prepare shape, strides, size info, but more importantly the offset-related data to loop through the outer, non-matrix dims
+            // Self and other's offset data
+            const selfOffsetShape = selfShape.slice(0, -2);
+            const otherOffsetShape = otherShape.slice(0, -2);
+            const selfOffsetStrides = selfStrides.slice(0, -2);
+            const otherOffsetStrides = otherStrides.slice(0, -2);
+            // The output's offset data
+            const offsetShape = Tensor.broadcastShapes(selfOffsetShape, otherOffsetShape);
+            const offsetSize = Tensor.shapeToSize(offsetShape);
+            const offsetStrides = Tensor.getStrides(offsetShape);
+            // Output shape, strides, size, value
+            const outputShape = [...offsetShape, batchARows, batchBCols];
+            const outputStrides = Tensor.getStrides(outputShape);
+            const outputSize = Tensor.shapeToSize(outputShape);
+            const outputValue = new Array(outputSize).fill(0);
+            // Loop through outer dims and do matmul on two outer-most dims
+            for (let index = 0; index < offsetSize; index++) {
+                const coords = Tensor.indexToCoords(index, offsetStrides);
+                const offset = Tensor.coordsToIndex(coords, outputStrides.slice(0, -2));
+                const selfOffset = Tensor.coordsToUnbroadcastedIndex(coords, selfOffsetShape, selfOffsetStrides);
+                const otherOffset = Tensor.coordsToUnbroadcastedIndex(coords, otherOffsetShape, otherOffsetStrides);
+                for (let i = 0; i < batchARows; i++) {
+                    for (let j = 0; j < batchBCols; j++) {
+                        for (let k = 0; k < batchACols; k++) {
+                            const outputIdx = offset + i * outputStrides[lastDim - 1] + j * outputStrides[lastDim];
+                            const selfIdx = selfOffset + i * selfStrides[lastDim - 1] + k * selfStrides[lastDim];
+                            const otherIdx = otherOffset + k * otherStrides[lastDim - 1] + j * otherStrides[lastDim];
+                            outputValue[outputIdx] += batchA[selfIdx] * batchB[otherIdx];
+                        }
+                    }
+                }
+            }
+            const out = new Tensor(outputValue, { shape: outputShape, strides: outputStrides });
+            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 = () => {
+                    other = other;
+                    const outGrad = out.grad.withGrad(false);
+                    const selfNoGrad = self.withGrad(false);
+                    const otherNoGrad = other.withGrad(false);
+                    if (this.requiresGrad)
+                        Tensor.addGrad(this, outGrad.matmul(otherNoGrad.transpose(lastDim - 1, lastDim)));
+                    if (other.requiresGrad)
+                        Tensor.addGrad(other, selfNoGrad.transpose(lastDim - 1, lastDim).matmul(outGrad));
+                };
+            }
+            return out;
+        }
         throw new Error(`Shapes [${this.shape}] and [${other.shape}] are not supported`);
     }
     // Utility to create a new tensor filled with a number

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "catniff",
-  "version": "0.2.13",
-  "description": "A cute autograd engine for Javascript",
+  "version": "0.2.15",
+  "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"
   ],