catniff 0.8.14 → 0.8.16

package/README.md

@@ -143,9 +143,8 @@ All available APIs are in [`./src/`](./src/) if you want to dig deeper.
 
 * More general tensor ops.
 * More general neural net APIs.
-* GPU acceleration.
-* Comprehensive caching.
-* Bug fixes.
+* GPU acceleration, possibly through WebGPU, Libtorch bindings, or CUDA.
+* Proper optimization.
 * More detailed documentation.
 * Code refactoring.
 * Proper tests.

package/dist/core.js

@@ -2504,7 +2504,7 @@ class Tensor {
                 visited.add(node);
                 // Reset grad to zeros if specified
                 if (zeroGrad) {
-                    node.grad = Tensor.zerosLike(node);
+                    delete node.grad;
                 }
                 for (let child of node.children)
                     build(child);

package/dist/lrscheduler.d.ts

@@ -5,9 +5,33 @@ export declare class StepLR {
     gamma: number;
     lastEpoch: number;
     baseLR: number;
+    baseGroupLRs: number[];
     constructor(optimizer: OptimizerWithLR, stepSize: number, gamma?: number, lastEpoch?: number);
-    step(epoch?: number): void;
+    step(): void;
+}
+export declare class LinearLR {
+    optimizer: OptimizerWithLR;
+    startFactor: number;
+    endFactor: number;
+    totalIters: number;
+    lastEpoch: number;
+    baseLR: number;
+    baseGroupLRs: number[];
+    constructor(optimizer: OptimizerWithLR, startFactor?: number, endFactor?: number, totalIters?: number, lastEpoch?: number);
+    step(): void;
+}
+export declare class CosineAnnealingLR {
+    optimizer: OptimizerWithLR;
+    TMax: number;
+    etaMin: number;
+    lastEpoch: number;
+    baseLR: number;
+    baseGroupLRs: number[];
+    constructor(optimizer: OptimizerWithLR, TMax: number, etaMin?: number, lastEpoch?: number);
+    step(): void;
 }
 export declare const LRScheduler: {
     StepLR: typeof StepLR;
+    LinearLR: typeof LinearLR;
+    CosineAnnealingLR: typeof CosineAnnealingLR;
 };

package/dist/lrscheduler.js

@@ -1,31 +1,91 @@
 "use strict";
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.LRScheduler = exports.StepLR = void 0;
+exports.LRScheduler = exports.CosineAnnealingLR = exports.LinearLR = exports.StepLR = void 0;
 class StepLR {
     optimizer;
     stepSize;
     gamma;
     lastEpoch;
     baseLR;
+    baseGroupLRs;
     constructor(optimizer, stepSize, gamma = 0.1, lastEpoch = -1) {
         this.optimizer = optimizer;
         this.stepSize = stepSize;
         this.gamma = gamma;
         this.lastEpoch = lastEpoch;
-        this.baseLR = this.optimizer.lr;
+        this.baseLR = optimizer.lr;
+        this.baseGroupLRs = this.optimizer.paramGroups.map(paramGroup => paramGroup.lr ?? this.optimizer.lr);
     }
-    step(epoch) {
-        if (typeof epoch === "undefined") {
-            this.lastEpoch++;
-            epoch = this.lastEpoch;
+    step() {
+        this.lastEpoch++;
+        // Update LR of each group
+        for (let index = 0; index < this.baseGroupLRs.length; index++) {
+            this.optimizer.paramGroups[index].lr = this.baseGroupLRs[index] * this.gamma ** Math.floor(this.lastEpoch / this.stepSize);
         }
-        else {
-            this.lastEpoch = epoch;
-        }
-        this.optimizer.lr = this.baseLR * this.gamma ** Math.floor(epoch / this.stepSize);
+        // Update default LR
+        this.optimizer.lr = this.baseLR * this.gamma ** Math.floor(this.lastEpoch / this.stepSize);
     }
 }
 exports.StepLR = StepLR;
+class LinearLR {
+    optimizer;
+    startFactor;
+    endFactor;
+    totalIters;
+    lastEpoch;
+    baseLR;
+    baseGroupLRs;
+    constructor(optimizer, startFactor = 0.3333333333333333, endFactor = 1, totalIters = 5, lastEpoch = -1) {
+        this.optimizer = optimizer;
+        this.startFactor = startFactor;
+        this.endFactor = endFactor;
+        this.totalIters = totalIters;
+        this.lastEpoch = lastEpoch;
+        this.baseLR = optimizer.lr;
+        this.baseGroupLRs = this.optimizer.paramGroups.map(paramGroup => paramGroup.lr ?? this.optimizer.lr);
+    }
+    step() {
+        this.lastEpoch++;
+        // Clamp under total allowed iterations
+        const t = Math.min(this.lastEpoch, this.totalIters);
+        // Precalculate factor
+        const factor = this.startFactor + (t / this.totalIters) * (this.endFactor - this.startFactor);
+        // Update LR of each group
+        for (let index = 0; index < this.baseGroupLRs.length; index++) {
+            this.optimizer.paramGroups[index].lr = this.baseGroupLRs[index] * factor;
+        }
+        // Update default LR
+        this.optimizer.lr = this.baseLR * factor;
+    }
+}
+exports.LinearLR = LinearLR;
+class CosineAnnealingLR {
+    optimizer;
+    TMax;
+    etaMin;
+    lastEpoch;
+    baseLR;
+    baseGroupLRs;
+    constructor(optimizer, TMax, etaMin = 0, lastEpoch = -1) {
+        this.optimizer = optimizer;
+        this.TMax = TMax;
+        this.etaMin = etaMin;
+        this.lastEpoch = lastEpoch;
+        this.baseLR = optimizer.lr;
+        this.baseGroupLRs = this.optimizer.paramGroups.map(paramGroup => paramGroup.lr ?? this.optimizer.lr);
+    }
+    step() {
+        this.lastEpoch++;
+        const cosine = (1 + Math.cos((this.lastEpoch * Math.PI) / this.TMax)) / 2;
+        for (let index = 0; index < this.baseGroupLRs.length; index++) {
+            this.optimizer.paramGroups[index].lr = this.etaMin + (this.baseGroupLRs[index] - this.etaMin) * cosine;
+        }
+        this.optimizer.lr = this.etaMin + (this.baseLR - this.etaMin) * cosine;
+    }
+}
+exports.CosineAnnealingLR = CosineAnnealingLR;
 exports.LRScheduler = {
-    StepLR
+    StepLR,
+    LinearLR,
+    CosineAnnealingLR
 };

package/dist/optim.d.ts

@@ -1,11 +1,15 @@
 import { Tensor } from "./core";
-export interface BaseOptimizerOptions {
-    lr?: number;
+export interface BaseParamGroup {
+    params: Tensor[];
+    [key: string]: any;
 }
 export declare abstract class BaseOptimizer {
-    params: Tensor[];
-    constructor(params: Tensor[], options?: BaseOptimizerOptions);
-    zeroGrad(): void;
+    paramGroups: BaseParamGroup[];
+    constructor(params: Tensor[] | BaseParamGroup[]);
+    zeroGrad(del?: boolean): void;
+}
+export interface OptimizerWithLR extends BaseOptimizer {
+    lr: number;
 }
 export interface SGDOptions {
     lr?: number;
@@ -14,17 +18,18 @@ export interface SGDOptions {
     weightDecay?: number;
     nesterov?: boolean;
 }
-export interface OptimizerWithLR extends BaseOptimizer {
-    lr: number;
+export interface SGDParamGroup extends SGDOptions {
+    params: Tensor[];
 }
 export declare class SGD extends BaseOptimizer {
+    paramGroups: SGDParamGroup[];
     lr: number;
-    momentumBuffers: Map<Tensor, Tensor>;
     momentum: number;
     dampening: number;
     weightDecay: number;
     nesterov: boolean;
-    constructor(params: Tensor[], options?: SGDOptions);
+    momentumBuffers: Map<Tensor, Tensor>;
+    constructor(params: Tensor[] | SGDParamGroup[], options?: SGDOptions);
     step(): void;
 }
 export interface AdamOptions {
@@ -33,15 +38,19 @@ export interface AdamOptions {
     eps?: number;
     weightDecay?: number;
 }
+export interface AdamParamGroup extends AdamOptions {
+    params: Tensor[];
+}
 export declare class Adam extends BaseOptimizer {
+    paramGroups: AdamParamGroup[];
     lr: number;
-    momentumBuffers: Map<Tensor, Tensor>;
-    velocityBuffers: Map<Tensor, Tensor>;
-    stepCount: number;
     betas: [number, number];
     eps: number;
     weightDecay: number;
-    constructor(params: Tensor[], options?: AdamOptions);
+    momentumBuffers: Map<Tensor, Tensor>;
+    velocityBuffers: Map<Tensor, Tensor>;
+    stepCounts: Map<Tensor, number>;
+    constructor(params: Tensor[] | AdamParamGroup[], options?: AdamOptions);
     step(): void;
 }
 export interface AdamWOptions {
@@ -50,15 +59,19 @@ export interface AdamWOptions {
     eps?: number;
     weightDecay?: number;
 }
+export interface AdamWParamGroup extends AdamWOptions {
+    params: Tensor[];
+}
 export declare class AdamW extends BaseOptimizer {
+    paramGroups: AdamWParamGroup[];
     lr: number;
-    momentumBuffers: Map<Tensor, Tensor>;
-    velocityBuffers: Map<Tensor, Tensor>;
-    stepCount: number;
     betas: [number, number];
     eps: number;
     weightDecay: number;
-    constructor(params: Tensor[], options?: AdamWOptions);
+    momentumBuffers: Map<Tensor, Tensor>;
+    velocityBuffers: Map<Tensor, Tensor>;
+    stepCounts: Map<Tensor, number>;
+    constructor(params: Tensor[] | AdamWParamGroup[], options?: AdamWOptions);
     step(): void;
 }
 export declare const Optim: {

package/dist/optim.js

@@ -3,27 +3,39 @@ Object.defineProperty(exports, "__esModule", { value: true });
 exports.Optim = exports.AdamW = exports.Adam = exports.SGD = exports.BaseOptimizer = void 0;
 const core_1 = require("./core");
 class BaseOptimizer {
-    params;
-    constructor(params, options) {
-        this.params = params;
+    paramGroups;
+    constructor(params) {
+        if (params[0] instanceof core_1.Tensor) {
+            this.paramGroups = [{ params: params }];
+        }
+        else {
+            this.paramGroups = params;
+        }
     }
-    zeroGrad() {
-        for (let index = 0; index < this.params.length; index++) {
-            const param = this.params[index];
-            param.grad = core_1.Tensor.zerosLike(param);
+    zeroGrad(del = true) {
+        for (let index = 0; index < this.paramGroups.length; index++) {
+            const paramGroup = this.paramGroups[index];
+            for (const param of paramGroup.params) {
+                if (del) {
+                    delete param.grad;
+                }
+                else {
+                    param.grad = core_1.Tensor.zerosLike(param);
+                }
+            }
         }
     }
 }
 exports.BaseOptimizer = BaseOptimizer;
 class SGD extends BaseOptimizer {
     lr;
-    momentumBuffers = new Map();
     momentum;
     dampening;
     weightDecay;
     nesterov;
+    momentumBuffers = new Map();
     constructor(params, options) {
-        super(params, options);
+        super(params);
         this.lr = options?.lr ?? 0.001;
         this.momentum = options?.momentum ?? 0;
         this.dampening = options?.dampening ?? 0;
@@ -31,159 +43,180 @@ class SGD extends BaseOptimizer {
         this.nesterov = options?.nesterov ?? false;
     }
     step() {
-        for (const param of this.params) {
-            if (!param.grad || !param.requiresGrad)
-                continue;
-            let grad = param.grad.detach(), detachedParam = param.detach();
-            // Apply weight decay (L2 regularization)
-            if (this.weightDecay !== 0) {
-                grad = grad.add(detachedParam.mul(this.weightDecay));
-            }
-            // Apply momentum
-            if (this.momentum !== 0) {
-                let buf = this.momentumBuffers.get(param);
-                if (!buf) {
-                    // First time: initialize momentum buffer with current gradient
-                    buf = grad.clone();
-                    this.momentumBuffers.set(param, buf);
-                }
-                else {
-                    // Update momentum buffer: buf = momentum * buf + (1 - dampening) * grad
-                    buf = buf.mul(this.momentum).add(grad.mul(1 - this.dampening));
-                    this.momentumBuffers.set(param, buf);
+        for (const paramGroup of this.paramGroups) {
+            const lr = paramGroup.lr ?? this.lr;
+            const momentum = paramGroup.momentum ?? this.momentum;
+            const dampening = paramGroup.dampening ?? this.dampening;
+            const weightDecay = paramGroup.weightDecay ?? this.weightDecay;
+            const nesterov = paramGroup.nesterov ?? this.nesterov;
+            for (const param of paramGroup.params) {
+                if (!param.grad || !param.requiresGrad)
+                    continue;
+                let grad = param.grad.detach(), detachedParam = param.detach();
+                // Apply weight decay (L2 regularization)
+                if (weightDecay !== 0) {
+                    grad = grad.add(detachedParam.mul(weightDecay));
                 }
-                if (this.nesterov) {
-                    // Nesterov momentum: grad = grad + momentum * buf
-                    grad = grad.add(buf.mul(this.momentum));
-                }
-                else {
-                    // Standard momentum: use momentum buffer as gradient
-                    grad = buf;
+                // Apply momentum
+                if (momentum !== 0) {
+                    let buf = this.momentumBuffers.get(param);
+                    if (!buf) {
+                        // First time: initialize momentum buffer with current gradient
+                        buf = grad.clone();
+                        this.momentumBuffers.set(param, buf);
+                    }
+                    else {
+                        // Update momentum buffer: buf = momentum * buf + (1 - dampening) * grad
+                        buf = buf.mul(momentum).add(grad.mul(1 - dampening));
+                        this.momentumBuffers.set(param, buf);
+                    }
+                    if (nesterov) {
+                        // Nesterov momentum: grad = grad + momentum * buf
+                        grad = grad.add(buf.mul(momentum));
+                    }
+                    else {
+                        // Standard momentum: use momentum buffer as gradient
+                        grad = buf;
+                    }
                 }
+                // Update parameter: param = param - lr * grad
+                const newParam = detachedParam.sub(grad.mul(lr));
+                param.replace(newParam);
             }
-            // Update parameter: param = param - lr * grad
-            const newParam = detachedParam.sub(grad.mul(this.lr));
-            param.replace(newParam);
         }
     }
 }
 exports.SGD = SGD;
 class Adam extends BaseOptimizer {
     lr;
-    momentumBuffers = new Map(); // First moment (m_t)
-    velocityBuffers = new Map(); // Second moment (v_t)
-    stepCount = 0;
     betas;
     eps;
     weightDecay;
+    momentumBuffers = new Map(); // First moment (m_t)
+    velocityBuffers = new Map(); // Second moment (v_t)
+    stepCounts = new Map();
     constructor(params, options) {
-        super(params, options);
+        super(params);
         this.lr = options?.lr ?? 0.001;
         this.betas = options?.betas ?? [0.9, 0.999];
         this.eps = options?.eps ?? 1e-8;
         this.weightDecay = options?.weightDecay ?? 0;
     }
     step() {
-        this.stepCount++;
-        const beta1 = this.betas[0];
-        const beta2 = this.betas[1];
-        // Bias correction factors
-        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 || !param.requiresGrad)
-                continue;
-            let grad = param.grad.detach(), detachedParam = param.detach();
-            // Apply weight decay (L2 regularization)
-            if (this.weightDecay !== 0) {
-                grad = grad.add(detachedParam.mul(this.weightDecay));
-            }
-            // Get or initialize first moment buffer (momentum)
-            let momentumBuffer = this.momentumBuffers.get(param);
-            if (!momentumBuffer) {
-                momentumBuffer = core_1.Tensor.zerosLike(grad); // Initialize with zeros (same shape as grad)
+        for (const paramGroup of this.paramGroups) {
+            const lr = paramGroup.lr ?? this.lr;
+            const betas = paramGroup.betas ?? this.betas;
+            const eps = paramGroup.eps ?? this.eps;
+            const weightDecay = paramGroup.weightDecay ?? this.weightDecay;
+            for (const param of paramGroup.params) {
+                if (!param.grad || !param.requiresGrad)
+                    continue;
+                // Get current step for param, initialize if has not step before
+                const stepCount = (this.stepCounts.get(param) ?? 0) + 1;
+                this.stepCounts.set(param, stepCount);
+                // Bias correction factors
+                const [beta1, beta2] = betas;
+                const biasCorrection1 = 1 - Math.pow(beta1, stepCount);
+                const biasCorrection2 = 1 - Math.pow(beta2, stepCount);
+                let grad = param.grad.detach(), detachedParam = param.detach();
+                // Apply weight decay (L2 regularization)
+                if (weightDecay !== 0) {
+                    grad = grad.add(detachedParam.mul(weightDecay));
+                }
+                // Get or initialize first moment buffer (momentum)
+                let momentumBuffer = this.momentumBuffers.get(param);
+                if (!momentumBuffer) {
+                    momentumBuffer = core_1.Tensor.zerosLike(grad); // Initialize with zeros (same shape as grad)
+                    this.momentumBuffers.set(param, momentumBuffer);
+                }
+                // Get or initialize second moment buffer (velocity)
+                let velocityBuffer = this.velocityBuffers.get(param);
+                if (!velocityBuffer) {
+                    velocityBuffer = core_1.Tensor.zerosLike(grad); // Initialize with zeros (same shape as grad)
+                    this.velocityBuffers.set(param, velocityBuffer);
+                }
+                // Update biased first moment estimate: m_t = beta1 * m_{t-1} + (1 - beta1) * g_t
+                momentumBuffer = momentumBuffer.mul(beta1).add(grad.mul(1 - beta1));
                 this.momentumBuffers.set(param, momentumBuffer);
-            }
-            // Get or initialize second moment buffer (velocity)
-            let velocityBuffer = this.velocityBuffers.get(param);
-            if (!velocityBuffer) {
-                velocityBuffer = core_1.Tensor.zerosLike(grad); // Initialize with zeros (same shape as grad)
+                // Update biased second moment estimate: v_t = beta2 * v_{t-1} + (1 - beta2) * g_t^2
+                velocityBuffer = velocityBuffer.mul(beta2).add(grad.pow(2).mul(1 - beta2));
                 this.velocityBuffers.set(param, velocityBuffer);
+                // Compute bias-corrected first moment: m_hat_t = m_t / (1 - beta1^t)
+                const correctedMomentum = momentumBuffer.div(biasCorrection1);
+                // Compute bias-corrected second moment: v_hat_t = v_t / (1 - beta2^t)
+                const correctedVelocity = velocityBuffer.div(biasCorrection2);
+                // Update parameters: theta_t = theta_{t-1} - alpha * m_hat_t / (sqrt(v_hat_t) + epsilon)
+                const denom = correctedVelocity.sqrt().add(eps);
+                const stepSize = correctedMomentum.div(denom).mul(lr);
+                const newParam = detachedParam.sub(stepSize);
+                param.replace(newParam);
             }
-            // Update biased first moment estimate: m_t = beta1 * m_{t-1} + (1 - beta1) * g_t
-            momentumBuffer = momentumBuffer.mul(beta1).add(grad.mul(1 - beta1));
-            this.momentumBuffers.set(param, momentumBuffer);
-            // Update biased second moment estimate: v_t = beta2 * v_{t-1} + (1 - beta2) * g_t^2
-            velocityBuffer = velocityBuffer.mul(beta2).add(grad.pow(2).mul(1 - beta2));
-            this.velocityBuffers.set(param, velocityBuffer);
-            // Compute bias-corrected first moment: m_hat_t = m_t / (1 - beta1^t)
-            const correctedMomentum = momentumBuffer.div(biasCorrection1);
-            // Compute bias-corrected second moment: v_hat_t = v_t / (1 - beta2^t)
-            const correctedVelocity = velocityBuffer.div(biasCorrection2);
-            // Update parameters: theta_t = theta_{t-1} - alpha * m_hat_t / (sqrt(v_hat_t) + epsilon)
-            const denom = correctedVelocity.sqrt().add(this.eps);
-            const stepSize = correctedMomentum.div(denom).mul(this.lr);
-            const newParam = detachedParam.sub(stepSize);
-            param.replace(newParam);
         }
     }
 }
 exports.Adam = Adam;
 class AdamW extends BaseOptimizer {
     lr;
-    momentumBuffers = new Map(); // First moment (m_t)
-    velocityBuffers = new Map(); // Second moment (v_t)
-    stepCount = 0;
     betas;
     eps;
     weightDecay;
+    momentumBuffers = new Map(); // First moment (m_t)
+    velocityBuffers = new Map(); // Second moment (v_t)
+    stepCounts = new Map();
     constructor(params, options) {
-        super(params, options);
+        super(params);
         this.lr = options?.lr ?? 0.001;
         this.betas = options?.betas ?? [0.9, 0.999];
         this.eps = options?.eps ?? 1e-8;
         this.weightDecay = options?.weightDecay ?? 0.01;
     }
     step() {
-        this.stepCount++;
-        const beta1 = this.betas[0];
-        const beta2 = this.betas[1];
-        // Bias correction factors
-        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 || !param.requiresGrad)
-                continue;
-            let grad = param.grad.detach(), detachedParam = param.detach();
-            // Apply weight decay (L2 regularization)
-            detachedParam = detachedParam.sub(detachedParam.mul(this.weightDecay).mul(this.lr));
-            // Get or initialize first moment buffer (momentum)
-            let momentumBuffer = this.momentumBuffers.get(param);
-            if (!momentumBuffer) {
-                momentumBuffer = core_1.Tensor.zerosLike(grad); // Initialize with zeros (same shape as grad)
+        for (const paramGroup of this.paramGroups) {
+            const lr = paramGroup.lr ?? this.lr;
+            const betas = paramGroup.betas ?? this.betas;
+            const eps = paramGroup.eps ?? this.eps;
+            const weightDecay = paramGroup.weightDecay ?? this.weightDecay;
+            for (const param of paramGroup.params) {
+                if (!param.grad || !param.requiresGrad)
+                    continue;
+                // Get current step for param, initialize if has not step before
+                const stepCount = (this.stepCounts.get(param) ?? 0) + 1;
+                this.stepCounts.set(param, stepCount);
+                // Bias correction factors
+                const [beta1, beta2] = betas;
+                const biasCorrection1 = 1 - Math.pow(beta1, stepCount);
+                const biasCorrection2 = 1 - Math.pow(beta2, stepCount);
+                let grad = param.grad.detach(), detachedParam = param.detach();
+                // Apply weight decay (L2 regularization)
+                detachedParam = detachedParam.sub(detachedParam.mul(weightDecay).mul(lr));
+                // Get or initialize first moment buffer (momentum)
+                let momentumBuffer = this.momentumBuffers.get(param);
+                if (!momentumBuffer) {
+                    momentumBuffer = core_1.Tensor.zerosLike(grad); // Initialize with zeros (same shape as grad)
+                    this.momentumBuffers.set(param, momentumBuffer);
+                }
+                // Get or initialize second moment buffer (velocity)
+                let velocityBuffer = this.velocityBuffers.get(param);
+                if (!velocityBuffer) {
+                    velocityBuffer = core_1.Tensor.zerosLike(grad); // Initialize with zeros (same shape as grad)
+                    this.velocityBuffers.set(param, velocityBuffer);
+                }
+                // Update biased first moment estimate: m_t = beta1 * m_{t-1} + (1 - beta1) * g_t
+                momentumBuffer = momentumBuffer.mul(beta1).add(grad.mul(1 - beta1));
                 this.momentumBuffers.set(param, momentumBuffer);
-            }
-            // Get or initialize second moment buffer (velocity)
-            let velocityBuffer = this.velocityBuffers.get(param);
-            if (!velocityBuffer) {
-                velocityBuffer = core_1.Tensor.zerosLike(grad); // Initialize with zeros (same shape as grad)
+                // Update biased second moment estimate: v_t = beta2 * v_{t-1} + (1 - beta2) * g_t^2
+                velocityBuffer = velocityBuffer.mul(beta2).add(grad.pow(2).mul(1 - beta2));
                 this.velocityBuffers.set(param, velocityBuffer);
+                // Compute bias-corrected first moment: m_hat_t = m_t / (1 - beta1^t)
+                const correctedMomentum = momentumBuffer.div(biasCorrection1);
+                // Compute bias-corrected second moment: v_hat_t = v_t / (1 - beta2^t)
+                const correctedVelocity = velocityBuffer.div(biasCorrection2);
+                // Update parameters: theta_t = theta_{t-1} - alpha * m_hat_t / (sqrt(v_hat_t) + epsilon)
+                const denom = correctedVelocity.sqrt().add(eps);
+                const stepSize = correctedMomentum.div(denom).mul(lr);
+                const newParam = detachedParam.sub(stepSize);
+                param.replace(newParam);
             }
-            // Update biased first moment estimate: m_t = beta1 * m_{t-1} + (1 - beta1) * g_t
-            momentumBuffer = momentumBuffer.mul(beta1).add(grad.mul(1 - beta1));
-            this.momentumBuffers.set(param, momentumBuffer);
-            // Update biased second moment estimate: v_t = beta2 * v_{t-1} + (1 - beta2) * g_t^2
-            velocityBuffer = velocityBuffer.mul(beta2).add(grad.pow(2).mul(1 - beta2));
-            this.velocityBuffers.set(param, velocityBuffer);
-            // Compute bias-corrected first moment: m_hat_t = m_t / (1 - beta1^t)
-            const correctedMomentum = momentumBuffer.div(biasCorrection1);
-            // Compute bias-corrected second moment: v_hat_t = v_t / (1 - beta2^t)
-            const correctedVelocity = velocityBuffer.div(biasCorrection2);
-            // Update parameters: theta_t = theta_{t-1} - alpha * m_hat_t / (sqrt(v_hat_t) + epsilon)
-            const denom = correctedVelocity.sqrt().add(this.eps);
-            const stepSize = correctedMomentum.div(denom).mul(this.lr);
-            const newParam = detachedParam.sub(stepSize);
-            param.replace(newParam);
         }
     }
 }

package/package.json

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