catniff 0.5.0 → 0.5.2

package/README.md

@@ -76,6 +76,32 @@ optim.step();
 console.log("Updated weight:", w.data);  // Should move toward 3.0
 ```
 
+## Neural networks
+
+There are built-in neural network constructs in Catniff as well:
+```js
+const { Tensor, nn } = require("catniff");
+
+// Linear layer with input size of 20 and output size of 10
+const linear = nn.Linear(20, 10);
+// RNN cell with input size of 32 and hidden size of 64
+const rnnCell = nn.RNNCell(32, 64);
+// Same thing but using GRU
+const gruCell = nn.GRUCell(32, 64);
+// Same thing but using LSTM
+const lstmCell = nn.LSTMCell(32, 64);
+
+// Forward passes
+const a = Tensor.randn([20]);
+const b = Tensor.randn([32]);
+const c = Tensor.randn([64]);
+
+linear.forward(a);
+rnnCell.forward(b, c);
+gruCell.forward(b, c);
+lstmCell.forward(b, c, c);
+```
+
 And it can still do much more, check out the docs and examples below for more information.
 
 ## Documentation

package/dist/core.js

@@ -20,6 +20,13 @@ class Tensor {
         this.gradFn = options.gradFn || (() => { });
         this.children = options.children || [];
         this.device = options.device || "cpu";
+        // Move tensor to device
+        if (this.device !== "cpu") {
+            const backend = Tensor.backends.get(this.device);
+            if (backend && backend.transfer) {
+                backend.transfer(this);
+            }
+        }
     }
     // Utility to flatten an nD array to be 1D
     static flatten(tensor) {
@@ -1549,7 +1556,8 @@ class Tensor {
     to(device) {
         const backend = Tensor.backends.get(device);
         if (backend && backend.transfer) {
-            return backend.transfer(this);
+            backend.transfer(this);
+            return this;
         }
         throw new Error(`No device found to transfer tensor to or a handler is not implemented for device.`);
     }

package/dist/nn.d.ts

@@ -2,13 +2,64 @@ import { Tensor, TensorValue } from "./core";
 declare class Linear {
     weight: Tensor;
     bias?: Tensor;
-    constructor(inFeatures: number, outFeatures: number, bias?: boolean, customInit?: (shape: number[]) => Tensor);
+    constructor(inFeatures: number, outFeatures: number, bias?: boolean, device?: string);
     forward(input: Tensor | TensorValue): Tensor;
 }
+declare class RNNCell {
+    weightIH: Tensor;
+    weightHH: Tensor;
+    biasIH?: Tensor;
+    biasHH?: Tensor;
+    constructor(inputSize: number, hiddenSize: number, bias?: boolean, device?: string);
+    forward(input: Tensor | TensorValue, hidden: Tensor | TensorValue): Tensor;
+}
+declare class GRUCell {
+    weightIR: Tensor;
+    weightIZ: Tensor;
+    weightIN: Tensor;
+    weightHR: Tensor;
+    weightHZ: Tensor;
+    weightHN: Tensor;
+    biasIR?: Tensor;
+    biasIZ?: Tensor;
+    biasIN?: Tensor;
+    biasHR?: Tensor;
+    biasHZ?: Tensor;
+    biasHN?: Tensor;
+    constructor(inputSize: number, hiddenSize: number, bias?: boolean, device?: string);
+    forward(input: Tensor | TensorValue, hidden: Tensor | TensorValue): Tensor;
+}
+export declare class LSTMCell {
+    weightII: Tensor;
+    weightIF: Tensor;
+    weightIG: Tensor;
+    weightIO: Tensor;
+    weightHI: Tensor;
+    weightHF: Tensor;
+    weightHG: Tensor;
+    weightHO: Tensor;
+    biasII?: Tensor;
+    biasIF?: Tensor;
+    biasIG?: Tensor;
+    biasIO?: Tensor;
+    biasHI?: Tensor;
+    biasHF?: Tensor;
+    biasHG?: Tensor;
+    biasHO?: Tensor;
+    constructor(inputSize: number, hiddenSize: number, bias?: boolean, device?: string);
+    forward(input: Tensor | TensorValue, hidden: Tensor | TensorValue, cell: Tensor | TensorValue): [Tensor, Tensor];
+}
+interface StateDict {
+    [key: string]: any;
+}
 export declare const nn: {
     Linear: typeof Linear;
+    RNNCell: typeof RNNCell;
+    GRUCell: typeof GRUCell;
     state: {
         getParameters(model: any, visited?: WeakSet<object>): Tensor[];
+        getStateDict(model: any, prefix?: string, visited?: WeakSet<object>): StateDict;
+        loadStateDict(model: any, stateDict: StateDict, prefix?: string, visited?: WeakSet<object>): void;
     };
 };
 export {};

package/dist/nn.js

@@ -1,37 +1,154 @@
 "use strict";
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.nn = void 0;
+exports.nn = exports.LSTMCell = void 0;
 const core_1 = require("./core");
+function linearTransform(input, weight, bias) {
+    let output = input.matmul(weight.t());
+    if (bias) {
+        output = output.add(bias);
+    }
+    return output;
+}
 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]);
+    constructor(inFeatures, outFeatures, bias = true, device) {
+        const bound = 1 / Math.sqrt(inFeatures);
+        this.weight = core_1.Tensor.uniform([outFeatures, inFeatures], -bound, bound, { requiresGrad: true, device });
         if (bias) {
-            this.bias = initFunc([outFeatures]);
+            this.bias = core_1.Tensor.uniform([outFeatures], -bound, bound, { requiresGrad: true, device });
         }
     }
     forward(input) {
         input = core_1.Tensor.forceTensor(input);
-        let output = input.matmul(this.weight.t());
-        if (this.bias) {
-            output = output.add(this.bias);
+        return linearTransform(input, this.weight, this.bias);
+    }
+}
+function rnnTransform(input, hidden, inputWeight, hiddenWeight, inputBias, hiddenBias) {
+    let output = input.matmul(inputWeight.t()).add(hidden.matmul(hiddenWeight.t()));
+    if (inputBias) {
+        output = output.add(inputBias);
+    }
+    if (hiddenBias) {
+        output = output.add(hiddenBias);
+    }
+    return output;
+}
+class RNNCell {
+    weightIH;
+    weightHH;
+    biasIH;
+    biasHH;
+    constructor(inputSize, hiddenSize, bias = true, device) {
+        const bound = 1 / Math.sqrt(hiddenSize);
+        this.weightIH = core_1.Tensor.uniform([hiddenSize, inputSize], -bound, bound, { requiresGrad: true, device });
+        this.weightHH = core_1.Tensor.uniform([hiddenSize, hiddenSize], -bound, bound, { requiresGrad: true, device });
+        if (bias) {
+            this.biasIH = core_1.Tensor.uniform([hiddenSize], -bound, bound, { requiresGrad: true, device });
+            this.biasHH = core_1.Tensor.uniform([hiddenSize], -bound, bound, { requiresGrad: true, device });
+        }
+    }
+    forward(input, hidden) {
+        input = core_1.Tensor.forceTensor(input);
+        hidden = core_1.Tensor.forceTensor(hidden);
+        return rnnTransform(input, hidden, this.weightIH, this.weightHH, this.biasIH, this.biasHH).tanh();
+    }
+}
+class GRUCell {
+    weightIR;
+    weightIZ;
+    weightIN;
+    weightHR;
+    weightHZ;
+    weightHN;
+    biasIR;
+    biasIZ;
+    biasIN;
+    biasHR;
+    biasHZ;
+    biasHN;
+    constructor(inputSize, hiddenSize, bias = true, device) {
+        const bound = 1 / Math.sqrt(hiddenSize);
+        this.weightIR = core_1.Tensor.uniform([hiddenSize, inputSize], -bound, bound, { requiresGrad: true, device });
+        this.weightIZ = core_1.Tensor.uniform([hiddenSize, inputSize], -bound, bound, { requiresGrad: true, device });
+        this.weightIN = core_1.Tensor.uniform([hiddenSize, inputSize], -bound, bound, { requiresGrad: true, device });
+        this.weightHR = core_1.Tensor.uniform([hiddenSize, hiddenSize], -bound, bound, { requiresGrad: true, device });
+        this.weightHZ = core_1.Tensor.uniform([hiddenSize, hiddenSize], -bound, bound, { requiresGrad: true, device });
+        this.weightHN = core_1.Tensor.uniform([hiddenSize, hiddenSize], -bound, bound, { requiresGrad: true, device });
+        if (bias) {
+            this.biasIR = core_1.Tensor.uniform([hiddenSize], -bound, bound, { requiresGrad: true, device });
+            this.biasIZ = core_1.Tensor.uniform([hiddenSize], -bound, bound, { requiresGrad: true, device });
+            this.biasIN = core_1.Tensor.uniform([hiddenSize], -bound, bound, { requiresGrad: true, device });
+            this.biasHR = core_1.Tensor.uniform([hiddenSize], -bound, bound, { requiresGrad: true, device });
+            this.biasHZ = core_1.Tensor.uniform([hiddenSize], -bound, bound, { requiresGrad: true, device });
+            this.biasHN = core_1.Tensor.uniform([hiddenSize], -bound, bound, { requiresGrad: true, device });
         }
-        return output;
+    }
+    forward(input, hidden) {
+        input = core_1.Tensor.forceTensor(input);
+        hidden = core_1.Tensor.forceTensor(hidden);
+        const r = rnnTransform(input, hidden, this.weightIR, this.weightHR, this.biasIR, this.biasHR).sigmoid();
+        const z = rnnTransform(input, hidden, this.weightIZ, this.weightHZ, this.biasIZ, this.biasHZ).sigmoid();
+        const n = linearTransform(input, this.weightIN, this.biasIN).add(r.mul(linearTransform(hidden, this.weightHN, this.biasHN))).tanh();
+        return (z.neg().add(1).mul(n).add(z.mul(hidden)));
     }
 }
+class LSTMCell {
+    weightII;
+    weightIF;
+    weightIG;
+    weightIO;
+    weightHI;
+    weightHF;
+    weightHG;
+    weightHO;
+    biasII;
+    biasIF;
+    biasIG;
+    biasIO;
+    biasHI;
+    biasHF;
+    biasHG;
+    biasHO;
+    constructor(inputSize, hiddenSize, bias = true, device) {
+        const bound = 1 / Math.sqrt(hiddenSize);
+        this.weightII = core_1.Tensor.uniform([hiddenSize, inputSize], -bound, bound, { requiresGrad: true, device });
+        this.weightIF = core_1.Tensor.uniform([hiddenSize, inputSize], -bound, bound, { requiresGrad: true, device });
+        this.weightIG = core_1.Tensor.uniform([hiddenSize, inputSize], -bound, bound, { requiresGrad: true, device });
+        this.weightIO = core_1.Tensor.uniform([hiddenSize, inputSize], -bound, bound, { requiresGrad: true, device });
+        this.weightHI = core_1.Tensor.uniform([hiddenSize, hiddenSize], -bound, bound, { requiresGrad: true, device });
+        this.weightHF = core_1.Tensor.uniform([hiddenSize, hiddenSize], -bound, bound, { requiresGrad: true, device });
+        this.weightHG = core_1.Tensor.uniform([hiddenSize, hiddenSize], -bound, bound, { requiresGrad: true, device });
+        this.weightHO = core_1.Tensor.uniform([hiddenSize, hiddenSize], -bound, bound, { requiresGrad: true, device });
+        if (bias) {
+            this.biasII = core_1.Tensor.uniform([hiddenSize], -bound, bound, { requiresGrad: true, device });
+            this.biasIF = core_1.Tensor.uniform([hiddenSize], -bound, bound, { requiresGrad: true, device });
+            this.biasIG = core_1.Tensor.uniform([hiddenSize], -bound, bound, { requiresGrad: true, device });
+            this.biasIO = core_1.Tensor.uniform([hiddenSize], -bound, bound, { requiresGrad: true, device });
+            this.biasHI = core_1.Tensor.uniform([hiddenSize], -bound, bound, { requiresGrad: true, device });
+            this.biasHF = core_1.Tensor.uniform([hiddenSize], -bound, bound, { requiresGrad: true, device });
+            this.biasHG = core_1.Tensor.uniform([hiddenSize], -bound, bound, { requiresGrad: true, device });
+            this.biasHO = core_1.Tensor.uniform([hiddenSize], -bound, bound, { requiresGrad: true, device });
+        }
+    }
+    forward(input, hidden, cell) {
+        input = core_1.Tensor.forceTensor(input);
+        hidden = core_1.Tensor.forceTensor(hidden);
+        cell = core_1.Tensor.forceTensor(cell);
+        const i = rnnTransform(input, hidden, this.weightII, this.weightHI, this.biasII, this.biasHI).sigmoid();
+        const f = rnnTransform(input, hidden, this.weightIF, this.weightHF, this.biasIF, this.biasHF).sigmoid();
+        const g = rnnTransform(input, hidden, this.weightIG, this.weightHG, this.biasIG, this.biasHG).tanh();
+        const o = rnnTransform(input, hidden, this.weightIO, this.weightHO, this.biasIO, this.biasHO).sigmoid();
+        const c = f.mul(cell).add(i.mul(g));
+        const h = o.mul(c.tanh());
+        return [h, c];
+    }
+}
+exports.LSTMCell = LSTMCell;
 const state = {
     getParameters(model, visited = new WeakSet()) {
-        if (visited.has(model)) {
+        if (visited.has(model))
             return [];
-        }
         visited.add(model);
         const parameters = [];
         for (const key in model) {
@@ -46,9 +163,47 @@ const state = {
             }
         }
         return parameters;
+    },
+    getStateDict(model, prefix = "", visited = new WeakSet()) {
+        if (visited.has(model))
+            return {};
+        visited.add(model);
+        const stateDict = {};
+        for (const key in model) {
+            if (!model.hasOwnProperty(key))
+                continue;
+            const value = model[key];
+            const fullKey = prefix ? `${prefix}.${key}` : key;
+            if (value instanceof core_1.Tensor) {
+                stateDict[fullKey] = value.val();
+            }
+            else if (typeof value === "object" && value !== null) {
+                Object.assign(stateDict, this.getStateDict(value, fullKey, visited));
+            }
+        }
+        return stateDict;
+    },
+    loadStateDict(model, stateDict, prefix = "", visited = new WeakSet()) {
+        if (visited.has(model))
+            return;
+        visited.add(model);
+        for (const key in model) {
+            if (!model.hasOwnProperty(key))
+                continue;
+            const value = model[key];
+            const fullKey = prefix ? `${prefix}.${key}` : key;
+            if (value instanceof core_1.Tensor && stateDict[fullKey]) {
+                value.replace(new core_1.Tensor(stateDict[fullKey]));
+            }
+            else if (typeof value === "object" && value !== null) {
+                this.loadStateDict(value, stateDict, fullKey, visited);
+            }
+        }
     }
 };
 exports.nn = {
     Linear,
+    RNNCell,
+    GRUCell,
     state
 };

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "catniff",
-  "version": "0.5.0",
+  "version": "0.5.2",
   "description": "A small Torch-like deep learning framework for Javascript",
   "main": "index.js",
   "scripts": {