catniff 0.2.2 → 0.2.4

package/dist/core.d.ts

@@ -9,8 +9,8 @@ export interface TensorOptions {
 }
 export declare class Tensor {
     value: number[] | number;
-    shape: number[];
-    strides: number[];
+    readonly shape: number[];
+    readonly strides: number[];
     grad?: Tensor;
     requiresGrad: boolean;
     gradFn: Function;

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") {
@@ -148,7 +153,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)) {
@@ -241,12 +250,14 @@ class Tensor {
                 }
             }
         }
+        // Remove size-1 dims only
         const outShape = this.shape.filter((dim, i) => {
             const shouldSqueeze = dims.includes(i);
             if (shouldSqueeze && dim !== 1)
                 throw new Error(`Can not squeeze dim with size ${dim}`);
             return !shouldSqueeze;
         });
+        // Remove strides of size-1 dims
         const outStrides = this.strides.filter((stride, i) => !dims.includes(i));
         const outValue = outShape.length === 0 ? this.value[0] : this.value;
         const out = new Tensor(outValue, {
@@ -280,12 +291,12 @@ class Tensor {
         // New stride
         const newStrides = [...this.strides];
         let newDimStride;
-        if (dim === this.shape.length) {
-            // Inserting at the back: reuse last stride or 1
-            newDimStride = this.strides[this.strides.length - 1] ?? 1;
+        if (dim >= this.shape.length) {
+            // Inserting at the back: use 1
+            newDimStride = 1;
         }
         else {
-            // Inserting before dim: copy current stride at that dim
+            // Inserting before dim: use current stride * current shape
             newDimStride = this.strides[dim] * this.shape[dim];
         }
         newStrides.splice(dim, 0, newDimStride);
@@ -317,8 +328,8 @@ class Tensor {
         const originalSize = this.shape.reduce((a, b) => a * b, 1);
         let gradShape, gradStrides, gradValue = [];
         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++) {
@@ -533,7 +544,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];
@@ -585,6 +596,7 @@ class Tensor {
         }
         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);
@@ -637,6 +649,7 @@ class Tensor {
         }
         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);
@@ -656,7 +669,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
@@ -670,6 +683,7 @@ class Tensor {
         }
         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);
@@ -703,7 +717,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() {
@@ -713,7 +727,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);
@@ -726,7 +740,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;
@@ -748,11 +762,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.2",
+  "version": "0.2.4",
   "description": "A cute autograd engine for Javascript",
   "main": "index.js",
   "scripts": {