catniff 0.5.0 → 0.5.1

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) {
@@ -1527,6 +1534,7 @@ class Tensor {
         return new Tensor(typeof this.value === "number" ? this.value : [...this.value], {
             shape: this.shape,
             strides: this.strides,
+            device: this.device,
             requiresGrad: this.requiresGrad
         });
     }
@@ -1549,7 +1557,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,11 +2,20 @@ 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;
+}
 export declare const nn: {
     Linear: typeof Linear;
+    RNNCell: typeof RNNCell;
     state: {
         getParameters(model: any, visited?: WeakSet<object>): Tensor[];
     };

package/dist/nn.js

@@ -5,17 +5,11 @@ 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]);
+    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) {
@@ -27,6 +21,31 @@ class Linear {
         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);
+        let output = input.matmul(this.weightIH.t())
+            .add(hidden.matmul(this.weightHH.t()));
+        if (this.biasIH && this.biasHH) {
+            output = output.add(this.biasIH).add(this.biasHH);
+        }
+        return output.tanh();
+    }
+}
 const state = {
     getParameters(model, visited = new WeakSet()) {
         if (visited.has(model)) {
@@ -50,5 +69,6 @@ const state = {
 };
 exports.nn = {
     Linear,
+    RNNCell,
     state
 };

package/package.json

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