catniff 0.4.1 → 0.5.0

package/README.md

@@ -9,22 +9,6 @@ Install through npm:
 npm install catniff
 ```
 
-## Example
-
-Here is a little demo of a quadratic function:
-```js
-const { Tensor } = require("catniff");
-
-const x = new Tensor(2, { requiresGrad: true });
-const L = x.pow(2).add(x); // x^2 + x
-
-L.backward();
-
-console.log(x.grad.val()); // 5
-```
-
-View all examples in [`./examples`](./examples).
-
 ## Tensors
 
 Tensors in Catniff can be created by passing in a number or an nD array, and there are built-in methods that can be used to perform tensor arithmetic:
@@ -92,7 +76,7 @@ optim.step();
 console.log("Updated weight:", w.data);  // Should move toward 3.0
 ```
 
-And it can still do much more, check out the docs mentioned below for more information.
+And it can still do much more, check out the docs and examples below for more information.
 
 ## Documentation
 
@@ -100,12 +84,19 @@ Full documentation is available in [`./docs/documentation.md`](./docs/documentat
 
 All available APIs are in [`./src/`](./src/) if you want to dig deeper.
 
+## Examples
+
+* [Simple neural net for XOR calculation](./examples/xornet.js).
+* [Tensors](./examples/tensors.js).
+* [Optimizer](./examples/optim.js).
+* [Simple quadratic equation](./examples/quadratic.js).
+
 ## Todos
 
 * Bug fixes.
 * More tensor ops.
 * GPU acceleration.
-* Some general neural net APIs.
+* More general neural net APIs.
 * More detailed documentation.
 * Code refactoring.
 * Proper tests.

package/dist/core.d.ts

@@ -46,6 +46,8 @@ export declare class Tensor {
     mean(dims?: number[] | number, keepDims?: boolean): Tensor;
     max(dims?: number[] | number, keepDims?: boolean): Tensor;
     min(dims?: number[] | number, keepDims?: boolean): Tensor;
+    var(dims?: number[] | number, keepDims?: boolean): Tensor;
+    std(dims?: number[] | number, keepDims?: boolean): Tensor;
     softmax(dims?: number[] | number): Tensor;
     add(other: TensorValue | Tensor): Tensor;
     sub(other: TensorValue | Tensor): Tensor;

package/dist/core.js

@@ -339,14 +339,20 @@ class Tensor {
     sum(dims, keepDims = false) {
         if (typeof this.value === "number")
             return this;
-        if (typeof dims === "number") {
-            dims = [dims];
-        }
         if (typeof dims === "undefined") {
             dims = Array.from({ length: this.shape.length }, (_, index) => index);
         }
+        if (Array.isArray(dims)) {
+            // Sort in descending order
+            const sortedDims = dims.sort((a, b) => b - a);
+            let reducedThis = this;
+            for (let i = 0; i < sortedDims.length; i++) {
+                reducedThis = reducedThis.sum(sortedDims[i], true);
+            }
+            return keepDims ? reducedThis : reducedThis.squeeze(dims);
+        }
         // Dims that are reduced now have size-1
-        const outputShape = this.shape.map((dim, i) => dims.includes(i) ? 1 : dim);
+        const outputShape = this.shape.map((dim, i) => dims === i ? 1 : dim);
         const outputStrides = Tensor.getStrides(outputShape);
         const outputSize = Tensor.shapeToSize(outputShape);
         const outputValue = new Array(outputSize).fill(0);
@@ -363,7 +369,7 @@ class Tensor {
         for (let realFlatIndex = 0; realFlatIndex < originalSize; realFlatIndex++) {
             const coords = Tensor.indexToCoords(realFlatIndex, this.strides);
             // Force 0 on reduced axes to collapse into size-1 dims
-            const outCoords = coords.map((val, i) => dims.includes(i) ? 0 : val);
+            const outCoords = coords.map((val, i) => dims === i ? 0 : val);
             // Convert output coordinates to flat index
             const outFlatIndex = Tensor.coordsToIndex(outCoords, outputStrides);
             // Add into sum
@@ -392,14 +398,20 @@ class Tensor {
     prod(dims, keepDims = false) {
         if (typeof this.value === "number")
             return this;
-        if (typeof dims === "number") {
-            dims = [dims];
-        }
         if (typeof dims === "undefined") {
             dims = Array.from({ length: this.shape.length }, (_, index) => index);
         }
+        if (Array.isArray(dims)) {
+            // Sort in descending order
+            const sortedDims = dims.sort((a, b) => b - a);
+            let reducedThis = this;
+            for (let i = 0; i < sortedDims.length; i++) {
+                reducedThis = reducedThis.prod(sortedDims[i], true);
+            }
+            return keepDims ? reducedThis : reducedThis.squeeze(dims);
+        }
         // Dims that are reduced now have size-1
-        const outputShape = this.shape.map((dim, i) => dims.includes(i) ? 1 : dim);
+        const outputShape = this.shape.map((dim, i) => dims === i ? 1 : dim);
         const outputStrides = Tensor.getStrides(outputShape);
         const outputSize = Tensor.shapeToSize(outputShape);
         const outputValue = new Array(outputSize).fill(1);
@@ -408,7 +420,7 @@ class Tensor {
         for (let realFlatIndex = 0; realFlatIndex < originalSize; realFlatIndex++) {
             const coords = Tensor.indexToCoords(realFlatIndex, this.strides);
             // Force 0 on reduced axes to collapse into size-1 dims
-            const outCoords = coords.map((val, i) => dims.includes(i) ? 0 : val);
+            const outCoords = coords.map((val, i) => dims === i ? 0 : val);
             // Convert output coordinates to flat index
             const outFlatIndex = Tensor.coordsToIndex(outCoords, outputStrides);
             // Multiply into product
@@ -427,7 +439,7 @@ class Tensor {
                 for (let realFlatIndex = 0; realFlatIndex < originalSize; realFlatIndex++) {
                     const coords = Tensor.indexToCoords(realFlatIndex, this.strides);
                     // Force 0 on reduced axes to collapse into size-1 dims
-                    const outCoords = coords.map((val, i) => dims.includes(i) ? 0 : val);
+                    const outCoords = coords.map((val, i) => dims === i ? 0 : val);
                     // Convert output coordinates to flat index
                     const outFlatIndex = Tensor.coordsToIndex(outCoords, outputStrides);
                     // Grad is the product of other elements of the same axis, which is product of all els divided by the current value
@@ -443,14 +455,20 @@ class Tensor {
     mean(dims, keepDims = false) {
         if (typeof this.value === "number")
             return this;
-        if (typeof dims === "number") {
-            dims = [dims];
-        }
         if (typeof dims === "undefined") {
             dims = Array.from({ length: this.shape.length }, (_, index) => index);
         }
+        if (Array.isArray(dims)) {
+            // Sort in descending order
+            const sortedDims = dims.sort((a, b) => b - a);
+            let reducedThis = this;
+            for (let i = 0; i < sortedDims.length; i++) {
+                reducedThis = reducedThis.mean(sortedDims[i], true);
+            }
+            return keepDims ? reducedThis : reducedThis.squeeze(dims);
+        }
         // Dims that are reduced now have size-1
-        const outputShape = this.shape.map((dim, i) => dims.includes(i) ? 1 : dim);
+        const outputShape = this.shape.map((dim, i) => dims === i ? 1 : dim);
         const outputStrides = Tensor.getStrides(outputShape);
         const outputSize = Tensor.shapeToSize(outputShape);
         const outputValue = new Array(outputSize).fill(0);
@@ -460,7 +478,7 @@ class Tensor {
         for (let realFlatIndex = 0; realFlatIndex < originalSize; realFlatIndex++) {
             const coords = Tensor.indexToCoords(realFlatIndex, this.strides);
             // Force 0 on reduced axes to collapse into size-1 dims
-            const outCoords = coords.map((val, i) => dims.includes(i) ? 0 : val);
+            const outCoords = coords.map((val, i) => dims === i ? 0 : val);
             // Convert output coordinates to flat index
             const outFlatIndex = Tensor.coordsToIndex(outCoords, outputStrides);
             // Calculate sum and contributors to the sum
@@ -485,7 +503,7 @@ class Tensor {
                 for (let realFlatIndex = 0; realFlatIndex < originalSize; realFlatIndex++) {
                     const coords = Tensor.indexToCoords(realFlatIndex, this.strides);
                     // Force 0 on reduced axes to collapse into size-1 dims
-                    const outCoords = coords.map((val, i) => dims.includes(i) ? 0 : val);
+                    const outCoords = coords.map((val, i) => dims === i ? 0 : val);
                     // Convert output coordinates to flat index
                     const outFlatIndex = Tensor.coordsToIndex(outCoords, outputStrides);
                     // Mean = 1/n * (el1 + el2 + ... + eln) so grad = 1/n
@@ -501,14 +519,20 @@ class Tensor {
     max(dims, keepDims = false) {
         if (typeof this.value === "number")
             return this;
-        if (typeof dims === "number") {
-            dims = [dims];
-        }
         if (typeof dims === "undefined") {
             dims = Array.from({ length: this.shape.length }, (_, index) => index);
         }
+        if (Array.isArray(dims)) {
+            // Sort in descending order
+            const sortedDims = dims.sort((a, b) => b - a);
+            let reducedThis = this;
+            for (let i = 0; i < sortedDims.length; i++) {
+                reducedThis = reducedThis.max(sortedDims[i], true);
+            }
+            return keepDims ? reducedThis : reducedThis.squeeze(dims);
+        }
         // Dims that are reduced now have size-1
-        const outputShape = this.shape.map((dim, i) => dims.includes(i) ? 1 : dim);
+        const outputShape = this.shape.map((dim, i) => dims === i ? 1 : dim);
         const outputStrides = Tensor.getStrides(outputShape);
         const outputSize = Tensor.shapeToSize(outputShape);
         const outputValue = new Array(outputSize).fill(-Infinity);
@@ -517,7 +541,7 @@ class Tensor {
         for (let realFlatIndex = 0; realFlatIndex < originalSize; realFlatIndex++) {
             const coords = Tensor.indexToCoords(realFlatIndex, this.strides);
             // Force 0 on reduced axes to collapse into size-1 dims
-            const outCoords = coords.map((val, i) => dims.includes(i) ? 0 : val);
+            const outCoords = coords.map((val, i) => dims === i ? 0 : val);
             // Convert output coordinates to flat index
             const outFlatIndex = Tensor.coordsToIndex(outCoords, outputStrides);
             // Get max over time
@@ -535,14 +559,25 @@ class Tensor {
             out.children.push(this);
             out.gradFn = () => {
                 const gradShape = this.shape, gradStrides = this.strides, gradValue = new Array(originalSize).fill(0);
+                const shareCounts = new Array(outputSize).fill(0);
+                const originalValue = this.value;
+                for (let realFlatIndex = 0; realFlatIndex < originalSize; realFlatIndex++) {
+                    const coords = Tensor.indexToCoords(realFlatIndex, this.strides);
+                    // Force 0 on reduced axes to collapse into size-1 dims
+                    const outCoords = coords.map((val, i) => dims === i ? 0 : val);
+                    // Convert output coordinates to flat index
+                    const outFlatIndex = Tensor.coordsToIndex(outCoords, outputStrides);
+                    // We collect how many elements share the same max value first
+                    shareCounts[outFlatIndex] += outputValue[outFlatIndex] === originalValue[realFlatIndex] ? 1 : 0;
+                }
                 for (let realFlatIndex = 0; realFlatIndex < originalSize; realFlatIndex++) {
                     const coords = Tensor.indexToCoords(realFlatIndex, this.strides);
                     // Force 0 on reduced axes to collapse into size-1 dims
-                    const outCoords = coords.map((val, i) => dims.includes(i) ? 0 : val);
+                    const outCoords = coords.map((val, i) => dims === i ? 0 : val);
                     // Convert output coordinates to flat index
                     const outFlatIndex = Tensor.coordsToIndex(outCoords, outputStrides);
-                    // Calculate grad by checking if a positon holds a value equal to the max value
-                    gradValue[realFlatIndex] = outputValue[outFlatIndex] === this.value[realFlatIndex] ? 1 : 0;
+                    // Here we share the grad between the elements that share the same max value
+                    gradValue[realFlatIndex] = outputValue[outFlatIndex] === originalValue[realFlatIndex] ? 1 / shareCounts[outFlatIndex] : 0;
                 }
                 const localGrad = new Tensor(gradValue, { shape: gradShape, strides: gradStrides });
                 Tensor.addGrad(this, out.grad.withGrad(false).mul(localGrad));
@@ -554,14 +589,20 @@ class Tensor {
     min(dims, keepDims = false) {
         if (typeof this.value === "number")
             return this;
-        if (typeof dims === "number") {
-            dims = [dims];
-        }
         if (typeof dims === "undefined") {
             dims = Array.from({ length: this.shape.length }, (_, index) => index);
         }
+        if (Array.isArray(dims)) {
+            // Sort in descending order
+            const sortedDims = dims.sort((a, b) => b - a);
+            let reducedThis = this;
+            for (let i = 0; i < sortedDims.length; i++) {
+                reducedThis = reducedThis.min(sortedDims[i], true);
+            }
+            return keepDims ? reducedThis : reducedThis.squeeze(dims);
+        }
         // Dims that are reduced now have size-1
-        const outputShape = this.shape.map((dim, i) => dims.includes(i) ? 1 : dim);
+        const outputShape = this.shape.map((dim, i) => dims === i ? 1 : dim);
         const outputStrides = Tensor.getStrides(outputShape);
         const outputSize = Tensor.shapeToSize(outputShape);
         const outputValue = new Array(outputSize).fill(Infinity);
@@ -570,7 +611,7 @@ class Tensor {
         for (let realFlatIndex = 0; realFlatIndex < originalSize; realFlatIndex++) {
             const coords = Tensor.indexToCoords(realFlatIndex, this.strides);
             // Force 0 on reduced axes to collapse into size-1 dims
-            const outCoords = coords.map((val, i) => dims.includes(i) ? 0 : val);
+            const outCoords = coords.map((val, i) => dims === i ? 0 : val);
             // Convert output coordinates to flat index
             const outFlatIndex = Tensor.coordsToIndex(outCoords, outputStrides);
             // Get min over time
@@ -588,14 +629,25 @@ class Tensor {
             out.children.push(this);
             out.gradFn = () => {
                 const gradShape = this.shape, gradStrides = this.strides, gradValue = new Array(originalSize).fill(0);
+                const shareCounts = new Array(outputSize).fill(0);
+                const originalValue = this.value;
+                for (let realFlatIndex = 0; realFlatIndex < originalSize; realFlatIndex++) {
+                    const coords = Tensor.indexToCoords(realFlatIndex, this.strides);
+                    // Force 0 on reduced axes to collapse into size-1 dims
+                    const outCoords = coords.map((val, i) => dims === i ? 0 : val);
+                    // Convert output coordinates to flat index
+                    const outFlatIndex = Tensor.coordsToIndex(outCoords, outputStrides);
+                    // We collect how many elements share the same min value first
+                    shareCounts[outFlatIndex] += outputValue[outFlatIndex] === originalValue[realFlatIndex] ? 1 : 0;
+                }
                 for (let realFlatIndex = 0; realFlatIndex < originalSize; realFlatIndex++) {
                     const coords = Tensor.indexToCoords(realFlatIndex, this.strides);
                     // Force 0 on reduced axes to collapse into size-1 dims
-                    const outCoords = coords.map((val, i) => dims.includes(i) ? 0 : val);
+                    const outCoords = coords.map((val, i) => dims === i ? 0 : val);
                     // Convert output coordinates to flat index
                     const outFlatIndex = Tensor.coordsToIndex(outCoords, outputStrides);
-                    // Calculate grad by checking if a positon holds a value equal to the min value
-                    gradValue[realFlatIndex] = outputValue[outFlatIndex] === this.value[realFlatIndex] ? 1 : 0;
+                    // Here we share the grad between the elements that share the same min value
+                    gradValue[realFlatIndex] = outputValue[outFlatIndex] === originalValue[realFlatIndex] ? 1 / shareCounts[outFlatIndex] : 0;
                 }
                 const localGrad = new Tensor(gradValue, { shape: gradShape, strides: gradStrides });
                 Tensor.addGrad(this, out.grad.withGrad(false).mul(localGrad));
@@ -603,18 +655,34 @@ class Tensor {
         }
         return keepDims ? out : out.squeeze(dims);
     }
+    // Tensor variance reduction
+    var(dims, keepDims = false) {
+        const meanXSquared = this.square().mean(dims, keepDims);
+        const meanXSquaredExpanded = this.mean(dims, keepDims).square();
+        return meanXSquared.sub(meanXSquaredExpanded);
+    }
+    // Tensor standard deviation reduction
+    std(dims, keepDims = false) {
+        return this.var(dims, keepDims).sqrt();
+    }
     // Tensor product reduction
     softmax(dims) {
         if (typeof this.value === "number")
             return this;
-        if (typeof dims === "number") {
-            dims = [dims];
-        }
         if (typeof dims === "undefined") {
             dims = Array.from({ length: this.shape.length }, (_, index) => index);
         }
+        if (Array.isArray(dims)) {
+            // Sort in descending order
+            const sortedDims = dims.sort((a, b) => b - a);
+            let reducedThis = this;
+            for (let i = 0; i < sortedDims.length; i++) {
+                reducedThis = reducedThis.softmax(sortedDims[i]);
+            }
+            return reducedThis;
+        }
         // Dims that are reduced now have size-1
-        const expSumShape = this.shape.map((dim, i) => dims.includes(i) ? 1 : dim);
+        const expSumShape = this.shape.map((dim, i) => dims === i ? 1 : dim);
         const expSumStrides = Tensor.getStrides(expSumShape);
         const expSumSize = Tensor.shapeToSize(expSumShape);
         const expSumValue = new Array(expSumSize).fill(0);
@@ -626,7 +694,7 @@ class Tensor {
         for (let realFlatIndex = 0; realFlatIndex < outputSize; realFlatIndex++) {
             const coords = Tensor.indexToCoords(realFlatIndex, outputStrides);
             // Force 0 on reduced axes to collapse into size-1 dims
-            const expSumCoords = coords.map((val, i) => dims.includes(i) ? 0 : val);
+            const expSumCoords = coords.map((val, i) => dims === i ? 0 : val);
             // Convert exp sum coordinates to flat index
             const expSumFlatIndex = Tensor.coordsToIndex(expSumCoords, expSumStrides);
             // Add e^x to the sum cache
@@ -636,7 +704,7 @@ class Tensor {
         for (let realFlatIndex = 0; realFlatIndex < outputSize; realFlatIndex++) {
             const coords = Tensor.indexToCoords(realFlatIndex, outputStrides);
             // Force 0 on reduced axes to collapse into size-1 dims
-            const expSumCoords = coords.map((val, i) => dims.includes(i) ? 0 : val);
+            const expSumCoords = coords.map((val, i) => dims === i ? 0 : val);
             // Convert exp sum coordinates to flat index
             const expSumFlatIndex = Tensor.coordsToIndex(expSumCoords, expSumStrides);
             // Calculate e^xi / sum

package/dist/nn.d.ts

@@ -0,0 +1,14 @@
+import { Tensor, TensorValue } from "./core";
+declare class Linear {
+    weight: Tensor;
+    bias?: Tensor;
+    constructor(inFeatures: number, outFeatures: number, bias?: boolean, customInit?: (shape: number[]) => Tensor);
+    forward(input: Tensor | TensorValue): Tensor;
+}
+export declare const nn: {
+    Linear: typeof Linear;
+    state: {
+        getParameters(model: any, visited?: WeakSet<object>): Tensor[];
+    };
+};
+export {};

package/dist/nn.js

@@ -0,0 +1,54 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.nn = void 0;
+const core_1 = require("./core");
+class Linear {
+    weight;
+    bias;
+    constructor(inFeatures, outFeatures, bias = true, customInit) {
+        let initFunc = (shape) => {
+            const bound = 1 / Math.sqrt(inFeatures);
+            return core_1.Tensor.uniform(shape, -bound, bound, { requiresGrad: true });
+        };
+        if (customInit) {
+            initFunc = customInit;
+        }
+        this.weight = initFunc([outFeatures, inFeatures]);
+        if (bias) {
+            this.bias = initFunc([outFeatures]);
+        }
+    }
+    forward(input) {
+        input = core_1.Tensor.forceTensor(input);
+        let output = input.matmul(this.weight.t());
+        if (this.bias) {
+            output = output.add(this.bias);
+        }
+        return output;
+    }
+}
+const state = {
+    getParameters(model, visited = new WeakSet()) {
+        if (visited.has(model)) {
+            return [];
+        }
+        visited.add(model);
+        const parameters = [];
+        for (const key in model) {
+            if (!model.hasOwnProperty(key))
+                continue;
+            const value = model[key];
+            if (value instanceof core_1.Tensor) {
+                parameters.push(value);
+            }
+            else if (typeof value === "object" && value !== null) {
+                parameters.push(...state.getParameters(value, visited));
+            }
+        }
+        return parameters;
+    }
+};
+exports.nn = {
+    Linear,
+    state
+};

package/dist/optim.js

@@ -20,9 +20,8 @@ class SGD {
     }
     step() {
         for (const param of this.params) {
-            if (!param.grad) {
-                throw new Error("Can not apply SGD on empty grad");
-            }
+            if (!param.grad || !param.requiresGrad)
+                continue;
             let grad = param.grad.detach(), detachedParam = param.detach();
             // Apply weight decay (L2 regularization)
             if (this.weightDecay !== 0) {
@@ -80,9 +79,8 @@ class Adam {
         const biasCorrection1 = 1 - Math.pow(beta1, this.stepCount);
         const biasCorrection2 = 1 - Math.pow(beta2, this.stepCount);
         for (const param of this.params) {
-            if (!param.grad) {
-                throw new Error("Can not apply Adam on empty grad");
-            }
+            if (!param.grad || !param.requiresGrad)
+                continue;
             let grad = param.grad.detach(), detachedParam = param.detach();
             // Apply weight decay (L2 regularization)
             if (this.weightDecay !== 0) {

package/index.d.ts

@@ -1,2 +1,3 @@
 export * from "./dist/core";
 export * from "./dist/optim";
+export * from "./dist/nn";

package/index.js

@@ -1,4 +1,5 @@
 module.exports = {
     ...require("./dist/core"),
-    ...require("./dist/optim")
+    ...require("./dist/optim"),
+    ...require("./dist/nn")
 };

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "catniff",
-  "version": "0.4.1",
-  "description": "A small Torch-like deep learning framework for Javascript with tensor and autograd support",
+  "version": "0.5.0",
+  "description": "A small Torch-like deep learning framework for Javascript",
   "main": "index.js",
   "scripts": {
     "test": "echo \"Error: no test specified\" && exit 1"