mini-jstorch 1.5.0 → 1.7.0

package/Docs/About.md

@@ -0,0 +1,83 @@
+# Mini-JSTorch — Technical Information
+
+---
+
+## General Information
+
+- **Project Name:** Mini-JSTorch  
+- **Internal Name:** JST (JSTorch)
+
+> Note:  
+> Early versions of Mini-JSTorch do not strictly follow semantic versioning conventions  
+> (e.g. `0.0.1` for patches, `0.1.0` for minor releases, `1.0.0` for major releases).  
+> This inconsistency reflects the early learning and experimental phase of the project.
+
+---
+
+## 1. Engine Architecture Limitations (JSTorch Core)
+
+This section outlines the known structural weaknesses of the JSTorch engine.  
+Although the architecture may appear complex, it is currently sensitive and tightly coupled.
+
+### Identified Limitations
+
+- **High dependency on Utilities**  
+  Every core class depends directly on the Utilities module, which is defined at the top of the `jstorch.js` file. This creates strong coupling across the engine.
+
+- **Limited Tensor dimensionality**  
+  Tensor implementations currently support only two dimensions.  
+  Extending support to higher-dimensional tensors would require significant architectural changes due to the existing complexity.
+
+- **Uneven class complexity**  
+  New or recently modified classes often become significantly more complex than others, leading to inconsistency in maintainability and internal design balance.
+
+---
+
+## 2. Rationale Behind the `fu_` Utilities
+
+This section explains why the `fu_` utilities were introduced despite the existence of internal Utilities.
+
+### Issues with Internal Utilities
+
+- The Utilities defined at the beginning of `jstorch.js` are **internal engine helpers**, not intended for direct user interaction.
+
+- These Utilities are heavily reused across multiple core classes.  
+  Any modification to a utility function may trigger **cascading (domino) errors** throughout the engine due to tight dependencies.
+
+- Some utility functions intentionally diverge from standard or expected formulas.  
+  For example:
+  - Expected formula:  
+    `Param1 - Param4 * Param3`
+  - Internal Utilities implementation:  
+    `Param1 - Param2 * Param3 + Param4`
+
+  This behavior exists because internal Utilities are optimized for class-level computations, not for user-facing correctness or predictability.
+
+### Purpose of `fu_` Utilities
+
+The `fu_` utilities were designed to improve the **user experience** by providing:
+
+- Predictable and correct computational behavior
+- User-friendly and stable helper functions
+- Isolation from internal engine changes
+- Reduced risk of incorrect outputs and dependency-based cascading errors
+
+In short, `fu_` exists to ensure safety, clarity, and consistency for end users of Mini-JSTorch.
+
+---
+
+## 3. SJK (Shortcut JST Keywords) Reference
+
+This section lists commonly used abbreviations and keywords within the Mini-JSTorch ecosystem.
+
+**Format:** `"KEYWORD" : "Full Name / Meaning"`
+
+- `"JST"` : JSTorch  
+- `"fu"` : For User / User-Friendly  
+- `"fun"` : Function  
+- `"Dummy"` : Experimental  
+- `"Exp"` : Restricted experimental entity  
+- `"msg"` : Message, comment, warning, announcement  
+- `"donot"` : Do not / Don't  
+
+---

package/Docs/Structure.md

@@ -0,0 +1,129 @@
+# Project File Structure # 
+
+This document describes the directory and file structure of the **mini-JSTorch** package.
+It provides an overview of how the project is organized and the purpose of each major component.
+
+---
+
+## Repository Overview
+
+```text
+mini-jstorch/
+├── demo/
+│   ├── fu_fun.js
+│   ├── MakeModel.js
+│   └── scheduler.js
+├── Docs/
+│   ├── About.md
+│   └── Structure.md
+├── src/
+│   ├── jstorch.js
+│   └── Dummy/
+│       └── msg/
+├── index.js
+├── package.json
+└── README.md
+```
+
+---
+
+## Directory Descriptions
+
+`/demo`
+
+- Contains demonstration and testing files.
+
+  - Used for unit testing, quick system checks, and example usage
+  - Intended for users who prefer practical examples over reading full API documentation
+  - Allows testing features without writing extensive manual code
+
+`/Docs`
+
+- Contains detailed documentation related to the mini-JSTorch package.
+
+  - Provides deeper explanations of internal design and usage
+  - Intended for contributors and advanced users
+
+`/src`
+
+- Contains the source code of the JSTorch engine.
+
+  - Houses all core logic and internal implementations
+  - Modifications in this directory directly affect engine behavior
+
+`/src/Dummy`
+
+- Experimental and restricted directory.
+
+  - Used for experimental purposes and future development
+  - Files inside this directory may be unstable or incomplete
+  - Not intended for public or production use
+
+`/src/Dummy/msg`
+
+- Contains warning or message files.
+
+  - Indicates that files within the `Dummy` directory are restricted
+  - Serves as a notification mechanism for experimental or future-update-related content
+
+---
+
+## File Descriptions
+
+`/demo/fu_fun.js`
+
+- Purpose: Tests all user-facing (`fu_`) functions
+- Notes: Focuses on friendly and predictable helper utilities
+
+`/demo/MakeModel.js`
+
+- Purpose: Demonstrates creation of a simple model
+- Notes: Uses the `StepLR` scheduler as part of the example workflow
+
+`/demo/scheduler.js`
+
+- Purpose: Tests scheduler-related functionality
+- Notes: Intended to validate learning rate scheduling behavior
+
+`/Docs/About.md`
+
+- Purpose: Contains additional information about the mini-JSTorch package
+- Notes: May include background, design decisions, or non-API-related explanations
+
+`/Docs/Structure.md`
+
+- Purpose: Documents the repository file and folder structure
+- Notes: This file
+
+`/src/jstorch.js`
+
+- Purpose: Core engine implementation
+
+- Notes:
+
+  - Contains all JSTorch engine logic and functions
+  - Central file of the entire package
+  - Changes here have wide-ranging effects
+
+`index.js`
+
+- Purpose: Package entry point
+- Notes: Exposes public APIs and connects internal modules
+
+`package.json`
+
+- Purpose: Project configuration and metadata
+- Notes: Defines dependencies, scripts, and package information
+
+`README.md`
+
+- Purpose: Main documentation entry
+- Notes: Provides overview, installation instructions, and basic usage
+
+**Notes**
+
+- Experimental files may change or be restricted without notice
+- Users are encouraged to rely on public APIs and documented utilities
+- Internal structures are subject to refactoring as the project evolves
+
+---

package/README.md

@@ -4,17 +4,17 @@ A lightweight JavaScript neural network library for rapid frontend AI experiment
 
 ## Overview
 
-Mini-JSTorch is a high-performance, minimalist JavaScript library for building neural networks. It runs efficiently in both frontend and backend environments, including low-end devices. The library enables quick experimentation and learning in AI without compromising stability, accuracy, or training reliability.
+Mini-JSTorch is a high-performance, minimalist JavaScript library for building neural networks. It runs efficiently in Frontend environments, including low-end devices. The library enables quick experimentation and learning in AI without compromising stability, accuracy, or training reliability.
 
-This release, **version 1.5.0:** Adds user-friendly tensor functions, Flatten layer, and improved Conv2D operations and Modify Some Class
-For Architecture Compability.
+This release, **version 1.7.0:** We introduce **Softmax Layer**, **Tokenizer**, **AdamW Optimizer**, and enhanced NLP capabilities.
 
 ---
 
 ## New Features Highlights
 
-- **User-Friendly Tensor API:** New `fu_` functions (`fu_tensor`, `fu_add`, `fu_matmul`, etc.) with automatic validation and shape checking
-- **Flatten Layer:** Essential for connecting CNN architectures to dense layers.
+- **Softmax Layer:** Professional classification output with proper gradient computation
+- **Tokenizer:** Lightweight text preprocessing for NLP tasks
+- **AdamW Optimizer:** Modern optimizer with decoupled weight decay
 
 ---
 
@@ -23,12 +23,17 @@ For Architecture Compability.
 - **Layers:** Linear, Flatten, Conv2D  
 - **Activations:** ReLU, Sigmoid, Tanh, LeakyReLU, GELU, Mish, SiLU, ELU  
 - **Loss Functions:** MSELoss, CrossEntropyLoss
-- **Optimizers:** Adam, SGD  
-- **Schedulers:** StepLR, LambdaLR
+- **Optimizers:** Adam, SGD, LION, **AdamW**  
+- **Schedulers:** StepLR, LambdaLR, ReduceLROnPlateau
 - **Regularization:** Dropout, BatchNorm2D  
 - **Utilities:** zeros, randomMatrix, softmax, crossEntropy, dot, addMatrices, reshape, stack, flatten, eye, concat  
 - **Model Container:** Sequential (for stacking layers with forward/backward passes)  
 
+# Others
+
+- **Tokenizer**
+- **Softmax Layer**
+
 ---
 
 ## Installation
@@ -43,7 +48,7 @@ npm install mini-jstorch
 ## Quick Start Example
 
 ```javascript
-import { Sequential, Linear, ReLU, Sigmoid, CrossEntropyLoss, Adam, StepLR } from './jstorch.js';
+import { Sequential, Linear, ReLU, Sigmoid, CrossEntropyLoss, Adam, StepLR } from './src/jstorch.js';
 
 // Build model
 const model = new Sequential([
@@ -101,21 +106,23 @@ loadModel(model2, json);
 
 ## Demos & Testing
 
-Check the `tests/` directory for ready-to-run demos:
-- **tests/MakeModel.js:** Build and run a simple neural network.
-- **tests/scheduler.js:** Experiment with learning rate schedulers.
+Check the `demo/` directory for ready-to-run demos:
+- **demo/MakeModel.js:** Build and run a simple neural network.
+- **demo/scheduler.js:** Experiment with learning rate schedulers.
+- **demo/fu_fun.js:** Test all user-friendly (fu or For Users/Friendly Users) functions
 - Add your own scripts for quick prototyping!
 
 ```bash
-node tests/MakeModel.js
-node tests/scheduler.js
+node demo/MakeModel.js
+node demo/scheduler.js
+node demo/fu_fun.js
 ```
 
 ---
 
 ## Intended Use Cases
 
- - Rapid prototyping of neural networks in frontend and backend.
+ - Rapid prototyping of neural networks in frontend.
  - Learning and teaching foundational neural network concepts.
  - Experimentation on low-end devices or mobile browsers.
  - Lightweight AI projects without GPU dependency.
@@ -126,22 +133,4 @@ node tests/scheduler.js
 
 `MIT License`
 
-**Copyright (c) 2025 rizal-editors**
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.
+**Copyright (c) 2024 rizal-editors**

package/{tests → demo}/fu_fun.js

@@ -5,68 +5,68 @@ import {
 } from '../src/jstorch.js';
 
 function testAllFuFunctions() {
-    console.log("🧪 TESTING ALL FU_FUNCTIONS\n");
+    console.log("TESTING ALL FU_FUNCTIONS\n");
     
     // Test 1: fu_tensor
     console.log("1. fu_tensor");
     const t1 = fu_tensor([[1, 2], [3, 4]]);
-    console.log("✅", t1.data);
+    console.log("", t1.data);
     
     // Test 2: fu_add
     console.log("\n2. fu_add");
     const a = fu_tensor([[1, 2]]);
     const b = fu_tensor([[3, 4]]);
     const c = fu_add(a, b);
-    console.log("✅", a.data, "+", b.data, "=", c.data);
+    console.log("", a.data, "+", b.data, "=", c.data);
     
     // Test 3: fu_mul
     console.log("\n3. fu_mul");
     const d = fu_mul(a, b);
-    console.log("✅", a.data, "*", b.data, "=", d.data);
+    console.log("", a.data, "*", b.data, "=", d.data);
     
     // Test 4: fu_matmul
     console.log("\n4. fu_matmul");
     const e = fu_tensor([[1, 2]]);
     const f = fu_tensor([[3], [4]]);
     const g = fu_matmul(e, f);
-    console.log("✅ matmul =", g.data);
+    console.log("matmul =", g.data);
     
     // Test 5: fu_sum & fu_mean
     console.log("\n5. fu_sum & fu_mean");
     const h = fu_tensor([[1, 2], [3, 4]]);
     const sum = fu_sum(h);
     const mean = fu_mean(h);
-    console.log("✅ sum =", sum.data, "mean =", mean.data);
+    console.log("sum =", sum.data, "mean =", mean.data);
     
     // Test 6: fu_relu
     console.log("\n6. fu_relu");
     const i = fu_tensor([[-1, 0], [1, 2]]);
     const relu = fu_relu(i);
-    console.log("✅ relu =", relu.data);
+    console.log("relu =", relu.data);
     
     // Test 7: fu_sigmoid
     console.log("\n7. fu_sigmoid");
     const sigmoid = fu_sigmoid(i);
-    console.log("✅ sigmoid =", sigmoid.data);
+    console.log("sigmoid =", sigmoid.data);
     
     // Test 8: fu_tanh
     console.log("\n8. fu_tanh");
     const tanh = fu_tanh(i);
-    console.log("✅ tanh =", tanh.data);
+    console.log("tanh =", tanh.data);
     
     // Test 9: fu_softmax
     console.log("\n9. fu_softmax");
     const j = fu_tensor([[1, 2, 3]]);
     const softmax = fu_softmax(j);
-    console.log("✅ softmax =", softmax.data);
+    console.log("softmax =", softmax.data);
     
     // Test 10: fu_flatten & fu_reshape
     console.log("\n10. fu_flatten & fu_reshape");
     const k = fu_tensor([[1, 2], [3, 4]]);
     const flat = fu_flatten(k);
     const reshaped = fu_reshape(flat, 1, 4);
-    console.log("✅ flatten =", flat.data);
-    console.log("✅ reshape =", reshaped.data);
+    console.log("flatten =", flat.data);
+    console.log("reshape =", reshaped.data);
 }
 
 testAllFuFunctions();

package/demo/scheduler.js

@@ -0,0 +1,69 @@
+// Example: Test ALL learning rate schedulers in mini-jstorch with mini-jstorch optimizers
+
+import { SGD, StepLR, LambdaLR, ReduceLROnPlateau, Tensor } from "../src/jstorch.js";
+
+const param = { param: [[1, 2], [3, 4]], grad: [[0, 0], [0, 0]] };
+const optimizer = new SGD([param], 0.1);
+
+// --- Test StepLR ---
+console.log("Testing StepLR...");
+const stepScheduler = new StepLR(optimizer, 3, 0.5);
+for (let epoch = 1; epoch <= 10; epoch++) {
+    stepScheduler.step();
+    console.log(`Epoch ${epoch}: LR = ${optimizer.lr.toFixed(4)}`);
+}
+
+// --- Test LambdaLR ---
+console.log("\nTesting LambdaLR...");
+optimizer.lr = 0.1; // Reset LR
+const lambdaScheduler = new LambdaLR(optimizer, epoch => 1.0 / (1 + epoch));
+for (let epoch = 1; epoch <= 5; epoch++) {
+    lambdaScheduler.step();
+    console.log(`Epoch ${epoch}: LR = ${optimizer.lr.toFixed(4)}`);
+}
+
+// --- Test ReduceLROnPlateau ---
+console.log("\nTesting ReduceLROnPlateau...");
+optimizer.lr = 0.1; // Reset LR
+const plateauScheduler = new ReduceLROnPlateau(optimizer, {
+    patience: 2,
+    factor: 0.5,
+    min_lr: 0.01,
+    verbose: true
+});
+
+// Simulate training with plateauing loss
+const losses = [0.9, 0.8, 0.7, 0.69, 0.68, 0.68, 0.68, 0.67, 0.67, 0.67];
+console.log("Simulated training with plateauing loss:");
+for (let epoch = 0; epoch < losses.length; epoch++) {
+    plateauScheduler.step(losses[epoch]);
+    console.log(`Epoch ${epoch + 1}: Loss = ${losses[epoch].toFixed(3)}, LR = ${optimizer.lr.toFixed(4)}, Wait = ${plateauScheduler.wait}`);
+}
+
+// --- Test ReduceLROnPlateau with Cooldown ---
+console.log("\nTesting ReduceLROnPlateau with Cooldown...");
+optimizer.lr = 0.1; // Reset LR
+const plateauWithCooldown = new ReduceLROnPlateau(optimizer, {
+    patience: 2,
+    factor: 0.5,
+    min_lr: 0.01,
+    cooldown: 2,
+    verbose: true
+});
+
+// Simulate training with multiple plateaus
+const losses2 = [0.9, 0.9, 0.9, 0.9, 0.8, 0.8, 0.8, 0.8, 0.7, 0.7];
+console.log("Simulated training with cooldown:");
+for (let epoch = 0; epoch < losses2.length; epoch++) {
+    plateauWithCooldown.step(losses2[epoch]);
+    console.log(`Epoch ${epoch + 1}: Loss = ${losses2[epoch].toFixed(3)}, LR = ${optimizer.lr.toFixed(4)}, Wait = ${plateauWithCooldown.wait}, Cooldown = ${plateauWithCooldown.cooldown_counter}`);
+}
+
+// --- Summary ---
+console.log("\nSCHEDULER SUMMARY:");
+console.log(`StepLR: ${stepScheduler.last_epoch} epochs processed`);
+console.log(`LambdaLR: ${lambdaScheduler.last_epoch} epochs processed`);
+console.log(`ReduceLROnPlateau: ${plateauScheduler.num_reductions} LR reductions`);
+console.log(`ReduceLROnPlateau with Cooldown: ${plateauWithCooldown.num_reductions} LR reductions`);
+
+console.log("\nAll schedulers tested successfully!");

package/index.js

@@ -1,2 +1,3 @@
 // package root
+// * = all exports from src/jstorch.js
 export * from "./src/jstorch.js";

package/package.json

@@ -1,22 +1,24 @@
 {
   "name": "mini-jstorch",
-  "version": "1.5.0",
+  "version": "1.7.0",
   "type": "module", 
-  "description": "A lightweight JavaScript neural network library for rapid frontend AI experimentation on low-resource devices Inspired by PyTorch.",
+  "description": "A lightweight JavaScript neural network library for learning AI concepts and rapid Frontend experimentation. PyTorch-inspired, zero dependencies, perfect for educational use.",
   "main": "index.js",
   "keywords": [
     "neural-network",
     "javascript",
     "lightweight-torch",
     "lightweight",
-    "small",
+    "small-torch",
     "javascript-torch",
-    "ai",
+    "ai-model",
     "jstorch",
     "pytorch",
     "front-end",
     "machine-learning",
-    "mini"
+    "tiny-ml",
+    "frontend-ai",
+    "mini-neural-network"
   ],
   "author": "Rizal",
   "license": "MIT",

package/src/jstorch.js

@@ -1,9 +1,9 @@
 /*!
  * Project: mini-jstorch
- * File: MainEngine.js
- * Author: M. Rizal H. (Actual Author Name)
+ * File: jstorch.js
+ * Author: Rizal-editors
  * License: MIT
- * Copyright (C) 2025 M. Rizal H.
+ * Copyright (C) 2025 Rizal-editors
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -24,8 +24,7 @@
  * SOFTWARE.
  */
 
-
-// ---------------------- Utilities ----------------------
+// ---------------------- DONOT USE THESE (ENGINE INTERNALS) ----------------------
 export function zeros(rows, cols) { 
     return Array.from({length:rows},()=>Array(cols).fill(0)); 
 }
@@ -34,10 +33,14 @@ export function ones(rows, cols) {
     return Array.from({length:rows},()=>Array(cols).fill(1)); 
 }
 
-export function randomMatrix(rows, cols, scale=0.1){ 
-    return Array.from({length:rows},()=>
-        Array.from({length:cols},()=> (Math.random()*2-1)*scale)
-    ); 
+export function randomMatrix(rows, cols, scale=null){
+    // Auto-scale based on layer size (Xavier init)
+    if (scale === null){
+        scale = Math.sqrt(2.0 / (rows + cols));
+    }
+
+    return Array.from({length: rows}, () =>
+        Array.from({length: cols}, () => (Math.random() * 2 - 1) * scale));
 }
 
 export function transpose(matrix){ 
@@ -71,7 +74,7 @@ export function crossEntropy(pred,target){
     return -target.reduce((sum,t,i)=>sum+t*Math.log(pred[i]+eps),0); 
 }
 
-// ---------------------- USERS FRIENDLY UTILS ----------------
+// ---------------------- USERS FRIENDLY UTILS (USE THIS!) ----------------
 export function fu_tensor(data, requiresGrad = false) {
     if (!Array.isArray(data) || !Array.isArray(data[0])) {
         throw new Error("fu_tensor: Data must be 2D array");
@@ -242,35 +245,75 @@ export class Tensor {
 
 // ---------------------- Layers ----------------------
 export class Linear {
-    constructor(inputDim,outputDim){
-        this.W=randomMatrix(inputDim,outputDim);
-        this.b=Array(outputDim).fill(0);
-        this.gradW=zeros(inputDim,outputDim);
-        this.gradb=Array(outputDim).fill(0);
-        this.x=null;
+    constructor(inputDim, outputDim){
+        this.W = randomMatrix(inputDim, outputDim);
+        this.b = Array(outputDim).fill(0);
+        this.gradW = zeros(inputDim, outputDim);
+        this.gradb = Array(outputDim).fill(0);
+        this.x = null;
+        this.originalShape = null; // Track input shape
     }
 
     forward(x){
-        this.x=x;
-        const out=dot(x,this.W);
-        return out.map((row,i)=>row.map((v,j)=>v+this.b[j]));
+        // Handle both [batch, features] and [batch, 1, features]
+        this.originalShape = this._getShapeType(x);
+        
+        if (this.originalShape === '3d') {
+            // Convert from [batch, 1, features] to [batch, features]
+            this.x = x.map(sample => sample[0]);
+        } else {
+            // Already in [batch, features] format
+            this.x = x;
+        }
+        
+        const out = dot(this.x, this.W);
+        return out.map((row, i) => row.map((v, j) => v + this.b[j]));
     }
 
     backward(grad){
-        for(let i=0;i<this.W.length;i++) for(let j=0;j<this.W[0].length;j++)
-            this.gradW[i][j]=this.x.reduce((sum,row,k)=>sum+row[i]*grad[k][j],0);
-        for(let j=0;j<this.b.length;j++)
-            this.gradb[j]=grad.reduce((sum,row)=>sum+row[j],0);
-
-        const gradInput=zeros(this.x.length,this.W.length);
-        for(let i=0;i<this.x.length;i++)
-            for(let j=0;j<this.W.length;j++)
-                for(let k=0;k<this.W[0].length;k++)
-                    gradInput[i][j]+=grad[i][k]*this.W[j][k];
+        // Compute gradients
+        for(let i = 0; i < this.W.length; i++) {
+            for(let j = 0; j < this.W[0].length; j++) {
+                this.gradW[i][j] = this.x.reduce((sum, row, k) => sum + row[i] * grad[k][j], 0);
+            }
+        }
+        
+        for(let j = 0; j < this.b.length; j++) {
+            this.gradb[j] = grad.reduce((sum, row) => sum + row[j], 0);
+        }
+
+        const gradInput = zeros(this.x.length, this.W.length);
+        for(let i = 0; i < this.x.length; i++) {
+            for(let j = 0; j < this.W.length; j++) {
+                for(let k = 0; k < this.W[0].length; k++) {
+                    gradInput[i][j] += grad[i][k] * this.W[j][k];
+                }
+            }
+        }
+        
+        //Convert back to original shape if needed
+        if (this.originalShape === '3d') {
+            return gradInput.map(row => [row]); // Back to [batch, 1, features]
+        }
         return gradInput;
     }
 
-    parameters(){ return [ {param:this.W,grad:this.gradW}, {param:[this.b],grad:[this.gradb]} ]; }
+    _getShapeType(x) {
+        if (Array.isArray(x[0]) && Array.isArray(x[0][0]) && !Array.isArray(x[0][0][0])) {
+            return '3d'; // [batch, 1, features]
+        } else if (Array.isArray(x[0]) && !Array.isArray(x[0][0])) {
+            return '2d'; // [batch, features]  
+        } else {
+            throw new Error(`Unsupported input shape for Linear layer`);
+        }
+    }
+
+    parameters(){ 
+        return [ 
+            {param: this.W, grad: this.gradW}, 
+            {param: [this.b], grad: [this.gradb]} 
+        ]; 
+    }
 }
 
 export class Flatten {
@@ -509,38 +552,187 @@ export class Sequential {
 
 // ---------------------- Activations ----------------------
 export class ReLU{ 
-    constructor(){ this.out=null; } 
+    constructor(){ this.mask = null; this.originalShape = null; } 
     
     forward(x){ 
-        // Handle both [batch, features] and [batch, channels, height, width]
-        if (Array.isArray(x[0]) && Array.isArray(x[0][0]) && Array.isArray(x[0][0][0])) {
-            // [batch, channels, height, width]
-            this.out = x.map(batch => 
-                batch.flatMap(channel => 
-                    channel.flatMap(row => 
-                        row.map(v => Math.max(0, v))
-                    )
-                )
-            );
+        this.originalShape = this._getShapeType(x);
+        
+        if (this.originalShape === '3d') {
+            // Handle [batch, 1, features]
+            this.mask = x.map(sample => sample[0].map(v => v > 0));
+            return x.map(sample => [sample[0].map(v => Math.max(0, v))]);
         } else {
-            // [batch, features] - existing behavior
-            this.out = x.map(r => r.map(v => Math.max(0, v)));
+            // Handle [batch, features]
+            this.mask = x.map(row => row.map(v => v > 0));
+            return x.map(row => row.map(v => Math.max(0, v)));
         }
-        return this.out; 
     } 
     
     backward(grad){ 
-        // Gradient shape must match forward output shape
-        if (Array.isArray(grad[0]) && Array.isArray(grad[0][0])) {
-            // Standard [batch, features]
-            return grad.map((r, i) => r.map((v, j) => v * (this.out[i][j] > 0 ? 1 : 0)));
+        if (this.originalShape === '3d') {
+            return grad.map((sample, i) => 
+                [sample[0].map((v, j) => this.mask[i][j] ? v : 0)]
+            );
         } else {
-            // return as is
-            return grad;
+            return grad.map((row, i) => 
+                row.map((v, j) => this.mask[i][j] ? v : 0)
+            );
         }
-    } 
+    }
+    
+    _getShapeType(x) {
+        if (Array.isArray(x[0]) && Array.isArray(x[0][0]) && !Array.isArray(x[0][0][0])) {
+            return '3d';
+        } else if (Array.isArray(x[0]) && !Array.isArray(x[0][0])) {
+            return '2d';
+        } else {
+            throw new Error(`Unsupported input shape for ReLU`);
+        }
+    }
+}
+
+// ---------------------- Softmax ----------------------
+export class Softmax {
+    constructor(dim = -1) {
+        this.dim = dim;
+        this.output = null;
+        this.input = null;
+    }
+
+    forward(x) {
+        this.input = x;
+        
+        // x: [batch_size, num_classes]
+        this.output = x.map(row => {
+            const maxVal = Math.max(...row);
+            const exps = row.map(v => Math.exp(v - maxVal));
+            const sumExps = exps.reduce((a, b) => a + b, 0);
+            return exps.map(v => v / sumExps);
+        });
+        return this.output;
+    }
+
+    backward(grad) {
+        // grad: [batch_size, num_classes] - gradient from next layer
+        const batchSize = grad.length;
+        const numClasses = grad[0].length;
+        
+        const gradInput = zeros(batchSize, numClasses);
+        
+        for (let i = 0; i < batchSize; i++) {
+            const s = this.output[i]; // Softmax output for this sample
+            const gradOut = grad[i];  // Gradient from loss
+            
+            // Compute Jacobian matrix: J_ij = s_i * (δ_ij - s_j)
+            for (let j = 0; j < numClasses; j++) {
+                let sum = 0;
+                for (let k = 0; k < numClasses; k++) {
+                    // J[j][k] = s[j] * ((j === k ? 1 : 0) - s[k])
+                    const jacobian = s[j] * ((j === k ? 1 : 0) - s[k]);
+                    sum += jacobian * gradOut[k];
+                }
+                gradInput[i][j] = sum;
+            }
+        }
+        
+        return gradInput;
+    }
+
+    parameters() {
+        return []; // Softmax has no trainable parameters
+    }
 }
 
+// ---------------------- Tokenizer ----------------------
+export class Tokenizer {
+    constructor(vocabSize = 2000){
+        this.vocabSize = vocabSize;
+        this.wordToIndex = new Map();
+        this.indexToWord = new Map();
+        this.fitted = false;
+    }
+
+    fit(texts){
+        const wordCounts = new Map();
+
+        // Count word frequencies from all texts
+        texts.forEach(text => {
+            const words = this._preprocess(text);
+            words.forEach(word => {
+                wordCounts.set(word, (wordCounts.get(word) || 0) + 1);
+            });
+        });
+
+        // Sort by frequency and take top words
+        const sortedWords = [...wordCounts.entries()]
+            .sort((a, b) => a[1] - a[1])
+            .slice(0, this.vocabSize - 1); // Reverse 1 for unknown
+
+        // Build vocabulary
+        this.wordToIndex.clear();
+        this.indexToWord.clear();
+
+        // Add unk token
+        this.wordToIndex.set('<UNK>', 0);
+        this.indexToWord.set(0, '<UNK>');
+
+        // Add most frequent words
+        sortedWords.forEach(([word], index) =>{
+            this.wordToIndex.set(word, index + 1);
+            this.indexToWord.set(index + 1, word);
+        })
+
+        this.fitted = true;
+        return this;
+    }
+
+    tokenize(text){
+        if (!this.fitted) throw new Error("Tokenizer not fitted. Call fit() first.");
+
+        const words = this._preprocess(text);
+        return words.map(word => this.wordToIndex.get(word) || 0);
+    }
+
+    tokenizeBatch(texts, maxLength=null){
+        if (!this.fitted) throw new Error("Tokenizer not fitted. Call fit() first.");
+
+        return texts.map(text => {
+            const tokens = this.tokenize(text);
+
+            if (maxLength !== null){
+                // Pad or truncate to maxLength
+                if (tokens.length > maxLength){
+                    return tokens.slice(0, maxLength);
+                } else {
+                    return [...tokens, ...Array(maxLength - tokens.length).fill(0)];
+                }
+            }
+
+            return tokens;
+        })
+    }
+
+    detokenize(tokens){
+        return tokens.map(token => this.indexToWord.get(token) || '<UNK>').join(' ');
+    }
+
+    detokenizeBatch(tokenBatches){
+        return tokenBatches.map(tokens => this.detokenize(tokens));
+    }
+
+    getVocabSize(){
+        return this.wordToIndex.size;
+    }
+
+    _preprocess(text) {
+        return text.toLowerCase()
+                  .replace(/[^\w\s]/g, ' ')  // Remove punctuation
+                  .split(/\s+/)             // Split by whitespace
+                  .filter(word => word.length > 0); // Remove empty strings
+    }
+}
+
+// I'm too lazy to break lines here, so everything stays in one line
 export class Sigmoid{ constructor(){ this.out=null; } forward(x){ const fn=v=>1/(1+Math.exp(-v)); this.out=x.map(r=>r.map(fn)); return this.out; } backward(grad){ return grad.map((r,i)=>r.map((v,j)=>v*this.out[i][j]*(1-this.out[i][j]))); } }
 export class Tanh{ constructor(){ this.out=null; } forward(x){ this.out=x.map(r=>r.map(v=>Math.tanh(v))); return this.out; } backward(grad){ return grad.map((r,i)=>r.map((v,j)=>v*(1-this.out[i][j]**2))); } }
 export class LeakyReLU{ constructor(alpha=0.01){ this.alpha=alpha; this.out=null; } forward(x){ this.out=x.map(r=>r.map(v=>v>0?v:v*this.alpha)); return this.out; } backward(grad){ return grad.map((r,i)=>r.map((v,j)=>v*(this.out[i][j]>0?1:this.alpha))); } }
@@ -555,24 +747,236 @@ export class CrossEntropyLoss{ forward(pred,target){ this.pred=pred; this.target
 
 // ---------------------- Optimizers ----------------------
 export class Adam{
-    constructor(params,lr=0.001,b1=0.9,b2=0.999,eps=1e-8){
-        this.params=params; this.lr=lr; this.beta1=b1; this.beta2=b2; this.eps=eps;
-        this.m=params.map(p=>zeros(p.param.length,p.param[0].length||1));
-        this.v=params.map(p=>zeros(p.param.length,p.param[0].length||1));
-        this.t=0;
+    constructor(params, lr = 0.001, b1 = 0.9, b2 = 0.999, eps = 1e-8, max_grad_norm = 1.0){
+        // Handle both parameter styles: (params, lr) OR (params, {lr, ...})
+        if (typeof lr === 'object') {
+            // Options object provided
+            const options = lr;
+            this.lr = options.lr || 0.001;
+            this.beta1 = options.b1 || options.beta1 || 0.9;
+            this.beta2 = options.b2 || options.beta2 || 0.999;
+            this.eps = options.eps || 1e-8;
+            this.max_grad_norm = options.max_grad_norm || 1.0;
+        } else {
+            // Individual parameters provided
+            this.lr = lr;
+            this.beta1 = b1;
+            this.beta2 = b2;
+            this.eps = eps;
+            this.max_grad_norm = max_grad_norm;
+        }
+        
+        this.params = params;
+        this.m = params.map(p => zeros(p.param.length, p.param[0].length || 1));
+        this.v = params.map(p => zeros(p.param.length, p.param[0].length || 1));
+        this.t = 0;
+        
     }
+
     step(){
         this.t++;
-        this.params.forEach((p,idx)=>{
-            for(let i=0;i<p.param.length;i++)
-                for(let j=0;j<(p.param[0].length||1);j++){
-                    const g=p.grad[i][j];
-                    this.m[idx][i][j]=this.beta1*this.m[idx][i][j]+(1-this.beta1)*g;
-                    this.v[idx][i][j]=this.beta2*this.v[idx][i][j]+(1-this.beta2)*g*g;
-                    const mHat=this.m[idx][i][j]/(1-Math.pow(this.beta1,this.t));
-                    const vHat=this.v[idx][i][j]/(1-Math.pow(this.beta2,this.t));
-                    p.param[i][j]-=this.lr*mHat/(Math.sqrt(vHat)+this.eps);
+        this.params.forEach((p, idx) => {
+            // Calculate gradient norm for clipping
+            let grad_norm_sq = 0;
+            for (let i = 0; i < p.param.length; i++){
+                for (let j = 0; j < (p.param[0].length || 1); j++){
+                    const grad_val = p.grad[i] && p.grad[i][j] !== undefined ? p.grad[i][j] : 0;
+                    grad_norm_sq += grad_val * grad_val;
+                }
+            }
+
+            const grad_norm = Math.sqrt(grad_norm_sq);
+            const clip_scale = grad_norm > this.max_grad_norm ? this.max_grad_norm / grad_norm : 1.0;
+
+            // Update with clipped gradients
+            for (let i = 0; i < p.param.length; i++){
+                for(let j = 0; j < (p.param[0].length || 1); j++){
+                    if (p.grad[i] && p.grad[i][j] !== undefined){
+                        const g = p.grad[i][j] * clip_scale;
+                        this.m[idx][i][j] = this.beta1 * this.m[idx][i][j] + (1 - this.beta1) * g;
+                        this.v[idx][i][j] = this.beta2 * this.v[idx][i][j] + (1 - this.beta2) * g * g;
+                        const mHat = this.m[idx][i][j] / (1 - Math.pow(this.beta1, this.t));
+                        const vHat = this.v[idx][i][j] / (1 - Math.pow(this.beta2, this.t));
+                        p.param[i][j] -= this.lr * mHat / (Math.sqrt(vHat) + this.eps);
+                    }
+                }
+            }
+        });
+    }
+}
+
+// ---------------------- AdamW Optimizer ----------------------
+export class AdamW {
+    constructor(params, options = {}) {
+        const {
+            lr = 0.001,
+            beta1 = 0.9,
+            beta2 = 0.999,
+            eps = 1e-8,
+            weight_decay = 0.01,  
+            max_grad_norm = 1.0
+        } = options;
+        
+        this.params = params;
+        this.lr = lr;
+        this.beta1 = beta1;
+        this.beta2 = beta2;
+        this.eps = eps;
+        this.weight_decay = weight_decay;
+        this.max_grad_norm = max_grad_norm;
+        
+        this.m = params.map(p => zeros(p.param.length, p.param[0].length || 1));
+        this.v = params.map(p => zeros(p.param.length, p.param[0].length || 1));
+        this.t = 0;
+    }
+
+    step() {
+        this.t++;
+        this.params.forEach((p, idx) => {
+            // Gradient clipping (same as Adam)
+            let grad_norm_sq = 0;
+            for (let i = 0; i < p.param.length; i++) {
+                for (let j = 0; j < (p.param[0].length || 1); j++) {
+                    const grad_val = p.grad[i] && p.grad[i][j] !== undefined ? p.grad[i][j] : 0;
+                    grad_norm_sq += grad_val * grad_val;
+                }
+            }
+            const grad_norm = Math.sqrt(grad_norm_sq);
+            const clip_scale = grad_norm > this.max_grad_norm ? this.max_grad_norm / grad_norm : 1.0;
+
+            // AdamW update: weight decay applied separately
+            for (let i = 0; i < p.param.length; i++) {
+                for (let j = 0; j < (p.param[0].length || 1); j++) {
+                    if (p.grad[i] && p.grad[i][j] !== undefined) {
+                        const g = p.grad[i][j] * clip_scale;
+                        
+                        // Adam moments
+                        this.m[idx][i][j] = this.beta1 * this.m[idx][i][j] + (1 - this.beta1) * g;
+                        this.v[idx][i][j] = this.beta2 * this.v[idx][i][j] + (1 - this.beta2) * g * g;
+                        
+                        const mHat = this.m[idx][i][j] / (1 - Math.pow(this.beta1, this.t));
+                        const vHat = this.v[idx][i][j] / (1 - Math.pow(this.beta2, this.t));
+                        
+                        // AdamW key difference: weight decay applied to weights, not gradients
+                        p.param[i][j] -= this.lr * (
+                            mHat / (Math.sqrt(vHat) + this.eps) + 
+                            this.weight_decay * p.param[i][j]  // Decoupled weight decay
+                        );
+                    }
+                }
+            }
+        });
+    }
+}
+
+export class SGD{
+    constructor(params, lr = 0.01, max_grad_norm = 1.0) {
+        this.params = params;
+        this.lr = lr;
+        this.max_grad_norm = max_grad_norm; // Gradient Clipping
+    }
+
+    step() {
+        this.params.forEach(p => {
+            // Calculate gradient norm
+            let grad_norm_sq = 0;
+            let total_params = 0;
+
+            for (let i = 0; i < p.param.length; i++){
+                const row = p.param[i];
+                for (let j = 0; j < (row.length || 1); j++) {
+                    const grad_val = p.grad[i] && p.grad[i][j] !== undefined ? p.grad[i][j] : 0;
+                    grad_norm_sq += grad_val * grad_val;
+                    total_params++;
+                }
+            }
+
+            const grad_norm = Math.sqrt(grad_norm_sq);
+
+            // Apply gradient clipping if needed
+            const clip_scale = grad_norm > this.max_grad_norm ? this.max_grad_norm / grad_norm : 1.0;
+
+            // Update parameters with clipped gradients
+            for (let i = 0; i < p.param.length; i++){
+                const row = p.param[i];
+                for (let j = 0; j < (row.length || 1); j++) {
+                    if (p.grad[i] && p.grad[i][j] !== undefined){
+                        p.param[i][j] -= this.lr * (p.grad[i][j] * clip_scale);
+                    }
+                }
+            }
+        });
+    }
+}
+
+
+export class LION {
+    constructor(params, options = {}) {
+        this.params = params;
+        
+        const {
+            lr = 0.0001,      // Lions typically uses smaller LR
+            beta1 = 0.9,      // First moment decay
+            beta2 = 0.99,     // Second moment decay  
+            weight_decay = 0, // L2 regularization
+            eps = 1e-8        // Numerical stability
+        } = options;
+        
+        this.lr = lr;
+        this.beta1 = beta1;
+        this.beta2 = beta2;
+        this.weight_decay = weight_decay;
+        this.eps = eps;
+        
+        // Initialize momentums
+        this.m = params.map(p => zeros(p.param.length, p.param[0].length || 1));
+        this.t = 0;
+    }
+
+    step() {
+        this.t++;
+        
+        this.params.forEach((p, idx) => {
+            for (let i = 0; i < p.param.length; i++) {
+                for (let j = 0; j < (p.param[0].length || 1); j++) {
+                    if (p.grad[i] && p.grad[i][j] !== undefined) {
+                        const grad = p.grad[i][j];
+                        
+                        // Update momentum: m_t = β1 * m_{t-1} + (1 - β1) * g_t
+                        this.m[idx][i][j] = this.beta1 * this.m[idx][i][j] + (1 - this.beta1) * grad;
+                        
+                        // LIONS update: param = param - η * sign(m_t + β2 * g_t)
+                        const update_term = this.m[idx][i][j] + this.beta2 * grad;
+                        
+                        // Get sign with epsilon for stability
+                        let sign_val;
+                        if (update_term > this.eps) sign_val = 1;
+                        else if (update_term < -this.eps) sign_val = -1;
+                        else sign_val = 0;
+                        
+                        let update = sign_val * this.lr;
+                        
+                        // Add weight decay if specified
+                        if (this.weight_decay > 0) {
+                            update += this.weight_decay * this.lr * p.param[i][j];
+                        }
+                        
+                        p.param[i][j] -= update;
+                    }
+                }
+            }
+        });
+    }
+
+    zeroGrad() {
+        this.params.forEach(p => {
+            if (p.grad) {
+                for (let i = 0; i < p.grad.length; i++) {
+                    for (let j = 0; j < p.grad[i].length; j++) {
+                        p.grad[i][j] = 0;
+                    }
                 }
+            }
         });
     }
 }
@@ -619,6 +1023,89 @@ export class LambdaLR {
     }
 }
 
+// ---------------------- ReduceLROnPlateau Scheduler ----------------------
+export class ReduceLROnPlateau {
+    constructor(optimizer, options = {}) {
+        this.optimizer = optimizer;
+        
+        // Destructure with defaults
+        const {
+            patience = 10,
+            factor = 0.5, 
+            min_lr = 1e-6,
+            threshold = 1e-4,
+            cooldown = 0,
+            verbose = false
+        } = options;
+        
+        this.patience = patience;
+        this.factor = factor;
+        this.min_lr = min_lr;
+        this.threshold = threshold;
+        this.cooldown = cooldown;
+        this.verbose = verbose;
+        
+        // State tracking
+        this.bestLoss = Infinity;
+        this.wait = 0;
+        this.cooldown_counter = 0;
+        this.num_reductions = 0;
+    }
+
+    step(loss) {
+        // Handle cooldown
+        if (this.cooldown_counter > 0) {
+            this.cooldown_counter--;
+            return;
+        }
+
+        // Check if this is significant improvement (relative threshold)
+        const improvement_needed = this.bestLoss * (1 - this.threshold);
+        const is_better = loss < improvement_needed;
+        
+        if (is_better) {
+            // Significant improvement - reset
+            this.bestLoss = loss;
+            this.wait = 0;
+        } else {
+            // No significant improvement
+            this.wait += 1;
+        }
+        
+        // Check if we've waited long enough
+        if (this.wait >= this.patience) {
+            this._reduce_lr();
+            this.cooldown_counter = this.cooldown;
+            this.wait = 0;
+        }
+    }
+
+    _reduce_lr() {
+        const old_lr = this.optimizer.lr;
+        const new_lr = Math.max(old_lr * this.factor, this.min_lr);
+        
+        if (new_lr < old_lr) {
+            this.optimizer.lr = new_lr;
+            this.num_reductions++;
+            
+            if (this.verbose) {
+                console.log(`ReduceLROnPlateau: reducing LR from ${old_lr} to ${new_lr}`);
+            }
+        }
+    }
+
+    get_last_lr() {
+        return this.optimizer.lr;
+    }
+
+    reset() {
+        this.bestLoss = Infinity;
+        this.wait = 0;
+        this.cooldown_counter = 0;
+        this.num_reductions = 0;
+    }
+}
+
 // ---------------------- ELU Activation ----------------------
 export class ELU {
     constructor(alpha=1.0) {
@@ -708,7 +1195,6 @@ export class SiLU {
     }
 }
 
-export class SGD{ constructor(params,lr=0.01){ this.params=params; this.lr=lr; } step(){ this.params.forEach(p=>{ for(let i=0;i<p.param.length;i++) for(let j=0;j<(p.param[0].length||1);j++) p.param[i][j]-=this.lr*p.grad[i][j]; }); } }
 
 // ---------------------- BatchNorm2D ----------------------
 export class BatchNorm2d {

package/tests/scheduler.js

@@ -1,23 +0,0 @@
-// Example: Test learning rate schedulers (StepLR and LambdaLR) with mini-jstorch optimizers
-
-import { SGD, StepLR, LambdaLR, Tensor } from "../src/jstorch.js";
-
-const param = { param: [[1, 2], [3, 4]], grad: [[0, 0], [0, 0]] };
-const optimizer = new SGD([param], 0.1);
-
-// --- Test StepLR ---
-console.log("Testing StepLR...");
-const stepScheduler = new StepLR(optimizer, 3, 0.5);
-for (let epoch = 1; epoch <= 10; epoch++) {
-    stepScheduler.step();
-    console.log(`Epoch ${epoch}: LR = ${optimizer.lr.toFixed(4)}`);
-}
-
-// --- Test LambdaLR ---
-console.log("\nTesting LambdaLR...");
-optimizer.lr = 0.1; // Reset LR
-const lambdaScheduler = new LambdaLR(optimizer, epoch => 1.0 / (1 + epoch));
-for (let epoch = 1; epoch <= 5; epoch++) {
-    lambdaScheduler.step();
-    console.log(`Epoch ${epoch}: LR = ${optimizer.lr.toFixed(4)}`);
-}