catniff 0.2.1 → 0.2.3

package/dist/core.js

@@ -241,12 +241,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,13 +282,13 @@ class Tensor {
         // New stride
         const newStrides = [...this.strides];
         let newDimStride;
-        if (dim === 0) {
-            // Inserting at front: use product of all original dimensions  
-            newDimStride = this.shape.reduce((a, b) => a * b, 1) || 1;
+        if (dim >= this.shape.length) {
+            // Inserting at the back: use 1
+            newDimStride = 1;
         }
         else {
-            // Inserting elsewhere: use stride of previous dimension
-            newDimStride = this.strides[dim - 1];
+            // Inserting before dim: use current stride * current shape
+            newDimStride = this.strides[dim] * this.shape[dim];
         }
         newStrides.splice(dim, 0, newDimStride);
         const out = new Tensor(this.value, { shape: newShape, strides: newStrides });
@@ -585,6 +587,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 +640,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);
@@ -670,6 +674,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);

package/package.json

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