catniff 0.5.6 → 0.5.7

package/dist/nn.d.ts

@@ -29,7 +29,7 @@ declare class GRUCell {
     constructor(inputSize: number, hiddenSize: number, bias?: boolean, device?: string);
     forward(input: Tensor | TensorValue, hidden: Tensor | TensorValue): Tensor;
 }
-export declare class LSTMCell {
+declare class LSTMCell {
     weightII: Tensor;
     weightIF: Tensor;
     weightIG: Tensor;
@@ -49,6 +49,14 @@ export declare class LSTMCell {
     constructor(inputSize: number, hiddenSize: number, bias?: boolean, device?: string);
     forward(input: Tensor | TensorValue, hidden: Tensor | TensorValue, cell: Tensor | TensorValue): [Tensor, Tensor];
 }
+declare class LayerNorm {
+    weight?: Tensor;
+    bias?: Tensor;
+    eps: number;
+    normalizedShape: number[];
+    constructor(normalizedShape: number | number[], eps?: number, elementwiseAffine?: boolean, bias?: boolean, device?: string);
+    forward(input: Tensor | TensorValue): Tensor;
+}
 interface StateDict {
     [key: string]: any;
 }
@@ -56,6 +64,8 @@ export declare const nn: {
     Linear: typeof Linear;
     RNNCell: typeof RNNCell;
     GRUCell: typeof GRUCell;
+    LSTMCell: typeof LSTMCell;
+    LayerNorm: typeof LayerNorm;
     state: {
         getParameters(model: any, visited?: WeakSet<object>): Tensor[];
         getStateDict(model: any, prefix?: string, visited?: WeakSet<object>): StateDict;

package/dist/nn.js

@@ -1,6 +1,6 @@
 "use strict";
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.nn = exports.LSTMCell = void 0;
+exports.nn = void 0;
 const core_1 = require("./core");
 function linearTransform(input, weight, bias) {
     let output = input.matmul(weight.t());
@@ -144,7 +144,51 @@ class LSTMCell {
         return [h, c];
     }
 }
-exports.LSTMCell = LSTMCell;
+class LayerNorm {
+    weight;
+    bias;
+    eps;
+    normalizedShape;
+    constructor(normalizedShape, eps = 1e-5, elementwiseAffine = true, bias = true, device) {
+        this.eps = eps;
+        this.normalizedShape = Array.isArray(normalizedShape) ? normalizedShape : [normalizedShape];
+        if (this.normalizedShape.length === 0) {
+            throw new Error("Normalized shape cannot be empty");
+        }
+        if (elementwiseAffine) {
+            this.weight = core_1.Tensor.ones(this.normalizedShape, { requiresGrad: true, device });
+            if (bias) {
+                this.bias = core_1.Tensor.zeros(this.normalizedShape, { requiresGrad: true, device });
+            }
+        }
+    }
+    forward(input) {
+        input = core_1.Tensor.forceTensor(input);
+        // Normalize over the specified dimensions
+        const normalizedDims = this.normalizedShape.length;
+        const startDim = input.shape.length - normalizedDims;
+        if (startDim < 0) {
+            throw new Error("Input does not have enough dims to normalize");
+        }
+        const dims = [];
+        for (let i = 0; i < normalizedDims; i++) {
+            if (input.shape[startDim + i] !== this.normalizedShape[i]) {
+                throw new Error(`Shape mismatch at dim ${startDim + i}: expected ${this.normalizedShape[i]}, got ${input.shape[startDim + i]}`);
+            }
+            dims.push(startDim + i);
+        }
+        const mean = input.mean(dims, true);
+        const variance = input.sub(mean).pow(2).mean(dims, true);
+        let normalized = input.sub(mean).div(variance.add(this.eps).sqrt());
+        if (this.weight) {
+            normalized = normalized.mul(this.weight);
+        }
+        if (this.bias) {
+            normalized = normalized.add(this.bias);
+        }
+        return normalized;
+    }
+}
 const state = {
     getParameters(model, visited = new WeakSet()) {
         if (visited.has(model))
@@ -178,7 +222,7 @@ const state = {
                 stateDict[fullKey] = value.val();
             }
             else if (typeof value === "object" && value !== null) {
-                Object.assign(stateDict, this.getStateDict(value, fullKey, visited));
+                Object.assign(stateDict, state.getStateDict(value, fullKey, visited));
             }
         }
         return stateDict;
@@ -196,7 +240,7 @@ const state = {
                 value.replace(new core_1.Tensor(stateDict[fullKey], { device: value.device }));
             }
             else if (typeof value === "object" && value !== null) {
-                this.loadStateDict(value, stateDict, fullKey, visited);
+                state.loadStateDict(value, stateDict, fullKey, visited);
             }
         }
     }
@@ -205,5 +249,7 @@ exports.nn = {
     Linear,
     RNNCell,
     GRUCell,
+    LSTMCell,
+    LayerNorm,
     state
 };

package/package.json

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