catniff 0.2.3 → 0.2.5

package/dist/core.d.ts

@@ -1,7 +1,7 @@
 export type TensorValue = number | TensorValue[];
 export interface TensorOptions {
-    shape?: number[];
-    strides?: number[];
+    shape?: readonly number[];
+    strides?: readonly number[];
     grad?: Tensor;
     requiresGrad?: boolean;
     gradFn?: Function;
@@ -9,20 +9,27 @@ export interface TensorOptions {
 }
 export declare class Tensor {
     value: number[] | number;
-    shape: number[];
-    strides: number[];
+    readonly shape: readonly number[];
+    readonly strides: readonly number[];
     grad?: Tensor;
     requiresGrad: boolean;
     gradFn: Function;
     children: Tensor[];
     constructor(value: TensorValue, options?: TensorOptions);
     static flatten(tensor: TensorValue): number[] | number;
-    static getShape(tensor: TensorValue): number[];
-    static getStrides(shape: number[]): number[];
-    static padShape(stridesA: number[], stridesB: number[], shapeA: number[], shapeB: number[]): number[][];
-    static broadcastShapes(shapeA: number[], shapeB: number[]): number[];
-    static indexToCoords(index: number, shape: number[], strides: number[]): number[];
-    static coordsToIndex(coords: number[], shape: number[], strides: number[]): number;
+    static getShape(tensor: TensorValue): readonly number[];
+    static getStrides(shape: readonly number[]): readonly number[];
+    static padShape(stridesA: readonly number[], stridesB: readonly number[], shapeA: readonly number[], shapeB: readonly number[]): [
+        readonly number[],
+        readonly number[],
+        readonly number[],
+        readonly number[]
+    ];
+    static broadcastShapes(shapeA: readonly number[], shapeB: readonly number[]): readonly number[];
+    static indexToCoords(index: number, shape: readonly number[], strides: readonly number[]): number[];
+    static coordsToUnbroadcastedIndex(coords: number[], shape: readonly number[], strides: readonly number[]): number;
+    static coordsToIndex(coords: number[], strides: readonly number[]): number;
+    static shapeToSize(shape: readonly number[]): number;
     static elementWiseAB(tA: Tensor, tB: Tensor, op: (tA: number, tB: number) => number): Tensor;
     static elementWiseSelf(tA: Tensor, op: (tA: number) => number): Tensor;
     elementWiseABDAG(other: TensorValue | Tensor, op: (a: number, b: number) => number, thisGrad?: (self: Tensor, other: Tensor, outGrad: Tensor) => Tensor, otherGrad?: (self: Tensor, other: Tensor, outGrad: Tensor) => Tensor): Tensor;

package/dist/core.js

@@ -20,8 +20,13 @@ class Tensor {
     }
     // Utility to flatten an nD array to be 1D
     static flatten(tensor) {
+        // Handle scalar tensors
         if (typeof tensor === "number")
             return tensor;
+        // If value is already 1D, we just need to return the value ('s reference)
+        if (typeof tensor[0] === "number")
+            return tensor;
+        // Or else recursively traverse through the nD array to flatten
         const result = [];
         function traverse(arr) {
             if (typeof arr === "number") {
@@ -88,7 +93,7 @@ class Tensor {
         }
         return newShape;
     }
-    // Convert flat index to array of coordinates
+    // Utility to convert flat index to array of coordinates
     static indexToCoords(index, shape, strides) {
         const coords = new Array(shape.length);
         let remaining = index;
@@ -100,17 +105,33 @@ class Tensor {
         }
         return coords;
     }
-    // Convert array of coordinates to *unbroadcasted* flat index 
-    static coordsToIndex(coords, shape, strides) {
+    // Utility to convert array of coordinates to *unbroadcasted* flat index 
+    static coordsToUnbroadcastedIndex(coords, shape, strides) {
         let index = 0;
         for (let i = 0; i < coords.length; i++) {
-            const coord = coords[i];
             // Handle broadcasting
-            const actualCoord = shape[i] === 1 ? 0 : coord;
+            const actualCoord = shape[i] === 1 ? 0 : coords[i];
             index += actualCoord * strides[i];
         }
         return index;
     }
+    // Utility to convert array of coordinates to flat index 
+    static coordsToIndex(coords, strides) {
+        let index = 0;
+        for (let i = 0; i < coords.length; i++) {
+            index += coords[i] * strides[i];
+        }
+        return index;
+    }
+    // Utility to convert shape into 1D value array size
+    static shapeToSize(shape) {
+        let prod = 1;
+        for (let i = 0; i < shape.length; i++) {
+            prod *= shape[i];
+        }
+        return prod;
+    }
+    ;
     // Utility for binary (two operators involved) element-wise ops
     static elementWiseAB(tA, tB, op) {
         if (typeof tA.value === "number" && typeof tB.value === "number") {
@@ -127,15 +148,15 @@ class Tensor {
         const outputShape = Tensor.broadcastShapes(paddedAShape, paddedBShape);
         // Get other output info
         const outputStrides = Tensor.getStrides(outputShape);
-        const outputSize = outputShape.reduce((a, b) => a * b, 1);
+        const outputSize = Tensor.shapeToSize(outputShape);
         const outputValue = new Array(outputSize);
         for (let i = 0; i < outputSize; i++) {
             // Get coordinates from 1D index
             const coordsOutput = Tensor.indexToCoords(i, outputShape, outputStrides);
             // Convert the coordinates to 1D index of flattened A with respect to A's shape
-            const indexA = Tensor.coordsToIndex(coordsOutput, paddedAShape, paddedAStrides);
+            const indexA = Tensor.coordsToUnbroadcastedIndex(coordsOutput, paddedAShape, paddedAStrides);
             // Convert the coordinates to 1D index of flattened B with respect to B's shape
-            const indexB = Tensor.coordsToIndex(coordsOutput, paddedBShape, paddedBStrides);
+            const indexB = Tensor.coordsToUnbroadcastedIndex(coordsOutput, paddedBShape, paddedBStrides);
             // Calculate with op
             outputValue[i] = op(tA.value[indexA], tB.value[indexB]);
         }
@@ -148,7 +169,11 @@ class Tensor {
     static elementWiseSelf(tA, op) {
         if (typeof tA.value === "number")
             return new Tensor(op(tA.value));
-        return new Tensor(tA.value.map(el => op(el)), { shape: [...tA.shape], strides: [...tA.strides] });
+        const newValue = new Array(tA.value.length);
+        for (let index = 0; index < tA.value.length; index++) {
+            newValue[index] = op(tA.value[index]);
+        }
+        return new Tensor(newValue, { shape: tA.shape, strides: tA.strides });
     }
     // Utility to do element-wise operation and build a dag node with another tensor
     elementWiseABDAG(other, op, thisGrad = () => new Tensor(0), otherGrad = () => new Tensor(0)) {
@@ -312,15 +337,18 @@ class Tensor {
         if (typeof dims === "undefined") {
             dims = Array.from({ length: this.shape.length }, (_, index) => index);
         }
+        // Dims that are reduced now have size-1
         const outputShape = this.shape.map((dim, i) => dims.includes(i) ? 1 : dim);
         const outputStrides = Tensor.getStrides(outputShape);
-        const outputSize = outputShape.reduce((a, b) => a * b, 1);
+        const outputSize = Tensor.shapeToSize(outputShape);
         const outputValue = new Array(outputSize).fill(0);
-        const originalSize = this.shape.reduce((a, b) => a * b, 1);
+        const originalSize = Tensor.shapeToSize(this.shape);
+        // Gradient data
         let gradShape, gradStrides, gradValue = [];
+        // Allocate gradient data only when needed
         if (this.requiresGrad) {
-            gradShape = [...this.shape];
-            gradStrides = [...this.strides];
+            gradShape = this.shape;
+            gradStrides = this.strides;
             gradValue = new Array(originalSize).fill(0);
         }
         for (let index = 0; index < originalSize; index++) {
@@ -328,11 +356,11 @@ class Tensor {
             // Force 0 on reduced axes to collapse into size-1 dims
             const outCoords = coords.map((val, i) => dims.includes(i) ? 0 : val);
             // Convert output coordinates to flat index
-            const outFlatIndex = outCoords.reduce((acc, val, i) => acc + val * outputStrides[i], 0);
-            // Accumulate
-            const realFlatIndex = coords.reduce((acc, val, i) => acc + val * this.strides[i], 0);
+            const outFlatIndex = Tensor.coordsToIndex(outCoords, outputStrides);
+            // Accumulate, outFlatIndex should match multiple realFlatIndexes
+            const realFlatIndex = Tensor.coordsToIndex(coords, this.strides);
             outputValue[outFlatIndex] += this.value[realFlatIndex];
-            // Mark for gradient
+            // Mark for gradient if needed
             if (this.requiresGrad) {
                 (gradValue)[realFlatIndex] = 1;
             }
@@ -535,7 +563,7 @@ class Tensor {
         }
         // If same dimension, return copy
         if (dim1 === dim2) {
-            return new Tensor(this.value, { shape: [...this.shape], strides: [...this.strides] });
+            return new Tensor(this.value, { shape: this.shape, strides: this.strides });
         }
         // Create new shape and strides by swapping
         const newShape = [...this.shape];
@@ -569,6 +597,7 @@ class Tensor {
         if (this.shape.length !== 1 || other.shape.length !== 1) {
             throw new Error("Inputs are not 1D tensors");
         }
+        // Simple vector dot product
         const vectLen = this.shape[0];
         const vectA = this.value;
         const vectB = other.value;
@@ -606,6 +635,7 @@ class Tensor {
         if (this.shape.length !== 2 || other.shape.length !== 2) {
             throw new Error("Inputs are not matrices");
         }
+        // Simple matrix multiplication
         const matA = this.value;
         const matB = other.value;
         const matAStrides = this.strides;
@@ -618,11 +648,12 @@ class Tensor {
             throw new Error("Invalid matrices shape for multiplication");
         const matCShape = [matARows, matBCols];
         const matCStrides = Tensor.getStrides(matCShape);
-        const matCSize = matCShape.reduce((a, b) => a * b, 1);
+        const matCSize = Tensor.shapeToSize(matCShape);
         const matC = new Array(matCSize).fill(0);
         for (let i = 0; i < matARows; i++) {
             for (let j = 0; j < matBCols; j++) {
                 for (let k = 0; k < matACols; k++) {
+                    // Tensor values are 1D arrays so we have to get real index using strides
                     matC[i * matCStrides[0] + j * matCStrides[1]] +=
                         matA[i * matAStrides[0] + k * matAStrides[1]] *
                             matB[k * matBStrides[0] + j * matBStrides[1]];
@@ -660,7 +691,7 @@ class Tensor {
             throw new Error("Input is not a 2D and 1D tensor pair");
         }
         // MM with no grad
-        const thisMat = new Tensor(this.value, { shape: [...this.shape], strides: [...this.strides] });
+        const thisMat = new Tensor(this.value, { shape: this.shape, strides: this.strides });
         const otherMat = new Tensor(other.value, { shape: [other.shape[0], 1], strides: [other.strides[0], 1] });
         const out = thisMat.mm(otherMat).squeeze(1);
         // Handle grad with original tensors
@@ -708,7 +739,7 @@ class Tensor {
     static fullLike(tensor, num, options = {}) {
         if (typeof tensor.value === "number")
             return new Tensor(num, options);
-        return new Tensor(tensor.value.map(el => num), { shape: [...tensor.shape], strides: [...tensor.strides], ...options });
+        return new Tensor(new Array(tensor.value.length).fill(num), { shape: tensor.shape, strides: tensor.strides, ...options });
     }
     // Reverse-mode autodiff call
     backward() {
@@ -718,7 +749,7 @@ class Tensor {
         function build(node) {
             if (!visited.has(node) && node.requiresGrad) {
                 visited.add(node);
-                node.grad = Tensor.fullLike(node, 0);
+                node.grad = Tensor.fullLike(node, 0); // Reset grad with 0
                 for (let child of node.children)
                     build(child);
                 topo.push(node);
@@ -731,7 +762,7 @@ class Tensor {
             topo[index].gradFn();
         }
     }
-    // Returns the number/nD array form of tensor
+    // Returns the raw number/nD array form of tensor
     val() {
         if (typeof this.value === "number")
             return this.value;
@@ -753,11 +784,11 @@ class Tensor {
         }
         return buildNested(this.value, this.shape, this.strides);
     }
-    // Returns a copy of the tensor with gradient turned on/off
+    // Returns a copy of the tensor with gradient turned on/off and detaches from autograd
     withGrad(requiresGrad) {
         return new Tensor(this.value, {
-            shape: [...this.shape],
-            strides: [...this.strides],
+            shape: this.shape,
+            strides: this.strides,
             requiresGrad
         });
     }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "catniff",
-  "version": "0.2.3",
+  "version": "0.2.5",
   "description": "A cute autograd engine for Javascript",
   "main": "index.js",
   "scripts": {