catniff 0.8.4 → 0.8.6

package/dist/core.d.ts

@@ -224,6 +224,7 @@ export declare class Tensor {
         zeroGrad?: boolean;
     }): void;
     val(): TensorValue;
+    toString(): string;
     detach(): Tensor;
     clone(): Tensor;
     replace(other: Tensor | TensorValue): Tensor;

package/dist/core.js

@@ -2396,6 +2396,53 @@ class Tensor {
         }
         return buildNested(this.value, this.shape, this.strides, this.offset);
     }
+    // Returns the nicely Pytorch-like formatted string form
+    toString() {
+        const val = this.val();
+        // Format a single number (integers get trailing dot)
+        const formatNum = (n) => {
+            if (Number.isInteger(n) && Math.abs(n) < 1e8) {
+                return n.toFixed(0) + ".";
+            }
+            return n.toString();
+        };
+        // Handle scalar
+        if (typeof val === "number") {
+            return `tensor(${formatNum(val)})`;
+        }
+        // Collect all numbers to find max width for alignment
+        const collectNumbers = (v) => {
+            if (typeof v === "number")
+                return [v];
+            return v.flatMap(collectNumbers);
+        };
+        const allNumbers = collectNumbers(val);
+        const maxWidth = Math.max(...allNumbers.map((n) => formatNum(n).length));
+        const ndim = this.shape.length;
+        const baseIndent = "tensor(".length; // 7
+        const formatNested = (v, depth) => {
+            if (typeof v === "number") {
+                return formatNum(v).padStart(maxWidth);
+            }
+            const arr = v;
+            // Innermost dimension: format as single line [x, y, z]
+            if (arr.length > 0 && typeof arr[0] === "number") {
+                const elements = arr.map((x) => formatNum(x).padStart(maxWidth));
+                return `[${elements.join(", ")}]`;
+            }
+            // Number of blank lines between elements at this depth
+            // Deeper = fewer blank lines (0 between rows, 1 between 2D blocks, etc.)
+            const blankLines = Math.max(0, ndim - depth - 2);
+            const separator = ",\n" + "\n".repeat(blankLines);
+            const innerIndent = " ".repeat(baseIndent + depth + 1);
+            const formatted = arr.map((item, i) => {
+                const str = formatNested(item, depth + 1);
+                return i === 0 ? "[" + str : innerIndent + str;
+            });
+            return formatted.join(separator) + "]";
+        };
+        return "tensor(" + formatNested(val, 0) + ")";
+    }
     // Returns a view of the tensor with gradient turned off and detaches from autograd
     detach() {
         return new Tensor(this.value, {

package/dist/nn.d.ts

@@ -50,6 +50,20 @@ export declare class LSTMCell {
     constructor(inputSize: number, hiddenSize: number, bias?: boolean, device?: string, dtype?: dtype);
     forward(input: Tensor | TensorValue, hidden: Tensor | TensorValue, cell: Tensor | TensorValue): [Tensor, Tensor];
 }
+export declare class BatchNorm {
+    weight?: Tensor;
+    bias?: Tensor;
+    runningMean?: Tensor;
+    runningVar?: Tensor;
+    eps: number;
+    momentum: number;
+    numFeatures: number;
+    affine: boolean;
+    trackRunningStats: boolean;
+    numBatchesTracked: number;
+    constructor(numFeatures: number, eps?: number, momentum?: number, affine?: boolean, trackRunningStats?: boolean, device?: string, dtype?: dtype);
+    forward(input: Tensor): Tensor;
+}
 export declare class LayerNorm {
     weight?: Tensor;
     bias?: Tensor;
@@ -91,6 +105,7 @@ export declare const nn: {
     RNNCell: typeof RNNCell;
     GRUCell: typeof GRUCell;
     LSTMCell: typeof LSTMCell;
+    BatchNorm: typeof BatchNorm;
     LayerNorm: typeof LayerNorm;
     RMSNorm: typeof RMSNorm;
     Embedding: typeof Embedding;

package/dist/nn.js

@@ -1,6 +1,6 @@
 "use strict";
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.nn = exports.MultiheadAttention = exports.Embedding = exports.RMSNorm = exports.LayerNorm = exports.LSTMCell = exports.GRUCell = exports.RNNCell = exports.Linear = void 0;
+exports.nn = exports.MultiheadAttention = exports.Embedding = exports.RMSNorm = exports.LayerNorm = exports.BatchNorm = exports.LSTMCell = exports.GRUCell = exports.RNNCell = exports.Linear = void 0;
 exports.scaledDotProductAttention = scaledDotProductAttention;
 const core_1 = require("./core");
 function linearTransform(input, weight, bias) {
@@ -149,6 +149,83 @@ class LSTMCell {
     }
 }
 exports.LSTMCell = LSTMCell;
+class BatchNorm {
+    weight;
+    bias;
+    runningMean;
+    runningVar;
+    eps;
+    momentum;
+    numFeatures;
+    affine;
+    trackRunningStats;
+    numBatchesTracked;
+    constructor(numFeatures, eps = 1e-5, momentum = 0.1, affine = true, trackRunningStats = true, device, dtype) {
+        this.numFeatures = numFeatures;
+        this.eps = eps;
+        this.momentum = momentum;
+        this.affine = affine;
+        this.trackRunningStats = trackRunningStats;
+        this.numBatchesTracked = 0;
+        if (this.affine) {
+            this.weight = core_1.Tensor.ones([numFeatures], { requiresGrad: true, device, dtype });
+            this.bias = core_1.Tensor.zeros([numFeatures], { requiresGrad: true, device, dtype });
+        }
+        if (this.trackRunningStats) {
+            this.runningMean = core_1.Tensor.zeros([numFeatures], { requiresGrad: false, device, dtype });
+            this.runningVar = core_1.Tensor.ones([numFeatures], { requiresGrad: false, device, dtype });
+        }
+    }
+    forward(input) {
+        // Input shape: (N, C, ...) where C = numFeatures
+        // Normalize over batch dimension and spatial dimensions (if any)
+        if (input.shape.length < 2) {
+            throw new Error("Input must have at least 2 dimensions (batch, features)");
+        }
+        if (input.shape[1] !== this.numFeatures) {
+            throw new Error(`Expected ${this.numFeatures} features, got ${input.shape[1]}`);
+        }
+        let mean;
+        let variance;
+        if (core_1.Tensor.training || !this.trackRunningStats) {
+            // Training or trackRunningStats disabled - calculate mean and variance from scratch
+            // Calculate mean and variance over batch and spatial dimensions
+            // Keep only the channel dimension
+            const dims = [0, ...Array.from({ length: input.shape.length - 2 }, (_, i) => i + 2)];
+            mean = input.mean(dims, true);
+            variance = input.sub(mean).pow(2).mean(dims, true);
+            // Update running statistics if enabled and in training mode
+            if (this.trackRunningStats && core_1.Tensor.training) {
+                const exponentialAverageFactor = this.momentum;
+                this.runningMean = this.runningMean
+                    .mul(1 - exponentialAverageFactor)
+                    .add(mean.squeeze().mul(exponentialAverageFactor));
+                // Use unbiased variance for running estimate
+                const n = input.shape.reduce((acc, val, idx) => idx === 1 ? acc : acc * val, 1);
+                const unbiasingFactor = n / (n - 1);
+                this.runningVar = this.runningVar
+                    .mul(1 - exponentialAverageFactor)
+                    .add(variance.squeeze().mul(exponentialAverageFactor * unbiasingFactor));
+                this.numBatchesTracked++;
+            }
+        }
+        else {
+            // Inference with trackRunningStats enabled - use running statistics
+            mean = this.runningMean.reshape([1, this.numFeatures, ...Array(input.shape.length - 2).fill(1)]);
+            variance = this.runningVar.reshape([1, this.numFeatures, ...Array(input.shape.length - 2).fill(1)]);
+        }
+        // Normalize
+        let normalized = input.sub(mean).div(variance.add(this.eps).sqrt());
+        // Apply affine transformation
+        if (this.affine) {
+            const weightReshaped = this.weight.reshape([1, this.numFeatures, ...Array(input.shape.length - 2).fill(1)]);
+            const biasReshaped = this.bias.reshape([1, this.numFeatures, ...Array(input.shape.length - 2).fill(1)]);
+            normalized = normalized.mul(weightReshaped).add(biasReshaped);
+        }
+        return normalized;
+    }
+}
+exports.BatchNorm = BatchNorm;
 class LayerNorm {
     weight;
     bias;
@@ -383,6 +460,7 @@ exports.nn = {
     RNNCell,
     GRUCell,
     LSTMCell,
+    BatchNorm,
     LayerNorm,
     RMSNorm,
     Embedding,

package/package.json

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