mini-jstorch 2.0.0 → 2.0.2

package/Docs/API.md

@@ -0,0 +1,277 @@
+# API Reference
+
+## Model Container
+
+### Sequential
+
+```js
+new Sequential(layers: Layer[])
+```
+
+Container that chains layers sequentially.
+
+**Methods:**
+- forward(x) — Pass input through all layers
+- backward(grad) — Backpropagate gradient through all layers
+- parameters() — Returns [{param, grad}, ...] for all trainable parameters
+- zeroGrad() — Zero all parameter gradients
+- train() — Set all layers to training mode
+- eval() — Set all layers to evaluation mode
+- stateDict() — Get {layer_0.weight, layer_0.bias, ...}
+- loadStateDict(dict) — Load weights from state dict object
+- step(lr) — Apply SGD step to all layers directly
+
+---
+
+## Layers
+
+### Linear
+
+```js
+new Linear(inFeatures: number, outFeatures: number)
+```
+
+Fully connected layer. Weight shape: [inFeatures, outFeatures]. Bias shape: [1, outFeatures].
+
+### Conv2D (experimental)
+
+```js
+new Conv2D(inChannels: number, outChannels: number, kernelSize: number, stride?: number, padding?: number)
+```
+
+2D convolution layer. Input shape: [batch, channels, height, width].
+
+### Flatten
+
+```js
+new Flatten()
+```
+
+Flattens multi-dimensional input per sample. Preserves batch dimension.
+
+---
+
+## Activations
+
+All activations follow the same interface:
+- forward(x) — x: [batch, features], returns: [batch, features]
+- backward(grad) — grad: [batch, features], returns: [batch, features]
+
+| Class | Formula |
+|-------|---------|
+| ReLU() | max(0, x) |
+| Sigmoid() | 1 / (1 + exp(-x)) |
+| Tanh() | tanh(x) |
+| LeakyReLU(alpha?) | x > 0 ? x : alpha * x |
+| GELU() | 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) |
+| ELU(alpha?) | x > 0 ? x : alpha * (exp(x) - 1) |
+| Mish() | x * tanh(ln(1 + exp(x))) |
+| SiLU() | x * sigmoid(x) |
+| Softmax(dim?) | exp(x - max) / sum(exp(x - max)) |
+
+---
+
+## Loss Functions
+
+### MSELoss
+
+```js
+new MSELoss()
+```
+
+loss = mean((pred - target)^2)
+gradient = 2 * (pred - target) / batchSize
+
+### SoftmaxCrossEntropyLoss (recommended for classification)
+
+```js
+new SoftmaxCrossEntropyLoss()
+```
+
+Combines softmax + cross-entropy in a numerically stable way. Input: logits (not probabilities). Do NOT combine with a Softmax layer.
+
+### BCEWithLogitsLoss (recommended for binary classification)
+
+```js
+new BCEWithLogitsLoss()
+```
+
+Combines sigmoid + binary cross-entropy. Numerically stable. Input: logits (not probabilities). Do NOT combine with a Sigmoid layer.
+
+### CrossEntropyLoss (deprecated)
+
+```js
+new CrossEntropyLoss()
+```
+
+Use SoftmaxCrossEntropyLoss instead. Exists for backward compatibility.
+
+---
+
+## Optimizers
+
+### Adam (recommended)
+
+```js
+new Adam(parameters, options)
+// or
+new Adam(parameters, lr, beta1, beta2, eps, maxGradNorm)
+```
+
+Options object: { lr: 0.001, b1: 0.9, b2: 0.999, eps: 1e-8, max_grad_norm: 1.0 }
+
+### AdamW
+
+```js
+new AdamW(parameters, options)
+```
+
+Adam with decoupled weight decay. Options include weight_decay (default: 0.01).
+
+### SGD
+
+```js
+new SGD(parameters, lr?, maxGradNorm?)
+```
+
+### Lion
+
+```js
+new LION(parameters, options)
+```
+
+Memory-efficient optimizer.
+
+**All optimizers have:**
+- step() — Update parameters using accumulated gradients
+- zeroGrad() — Zero all parameter gradients
+
+---
+
+## Learning Rate Schedulers
+
+### StepLR
+
+```js
+new StepLR(optimizer, stepSize, gamma)
+```
+
+Multiplies LR by gamma every stepSize steps.
+
+### LambdaLR
+
+```js
+new LambdaLR(optimizer, fn)
+```
+
+Sets LR = baseLr * fn(epoch) on each step.
+
+### ReduceLROnPlateau
+
+```js
+new ReduceLROnPlateau(optimizer, options)
+```
+
+Reduces LR when loss stops improving. Options: patience, factor, min_lr, threshold, cooldown, verbose.
+
+---
+
+## Regularization
+
+### Dropout
+
+```js
+new Dropout(p?)
+```
+
+Randomly zeros p fraction of neurons during training. Default p = 0.5. Call layer.train() / layer.eval() to toggle.
+
+### BatchNorm2d (experimental)
+
+```js
+new BatchNorm2d(numFeatures, eps?, momentum?, affine?)
+```
+
+2D batch normalization. Input shape: [batch, channels, height, width].
+
+---
+
+## Tokenizer
+
+```js
+new Tokenizer(vocabSize?)
+```
+
+| Method | Description |
+|--------|-------------|
+| fit(texts) | Build vocabulary from text array |
+| transform(texts, maxLength?, padToMax?) | Convert texts to token indices |
+| fitTransform(texts, maxLength?, padToMax?) | Fit + transform in one call |
+| inverseTransform(tokens, skipPad?) | Convert token indices back to text |
+| getVocabulary() | Get vocabulary as string array |
+| getVocabSize() | Get vocabulary size |
+| getWordCounts() | Get word frequency map |
+| getMostCommon(n?) | Get top N most frequent words |
+
+---
+
+## Utilities
+
+zeros(rows, cols) — Matrix filled with 0
+ones(rows, cols) — Matrix filled with 1
+randomMatrix(rows, cols, scale?) — Random matrix (Xavier init if scale omitted)
+transpose(matrix) — Matrix transpose
+dot(a, b) — Matrix multiplication
+addMatrices(a, b) — Element-wise addition
+softmax(vector) — Softmax on 1D array
+crossEntropy(pred, target) — Cross-entropy loss (scalar)
+reshape(tensor, rows, cols) — Reshape to new dimensions
+flattenBatch(batch) — Flatten batch to 2D
+concat(a, b, axis) — Concatenate along axis 0 or 1
+stack(tensors) — Stack tensors
+
+---
+
+## Model Persistence
+
+saveModel(model) — Serialize model to JSON string
+loadModel(model, json) — Load weights from JSON into model
+
+Supports all layer types. Validates layer types and shapes. Logs warnings for mismatches.
+
+---
+
+## Tensor (Advanced)
+
+```js
+new Tensor(data, requiresGrad?)
+```
+
+| Method | Description |
+|--------|-------------|
+| add(tensor) | Element-wise addition |
+| mul(tensor) | Element-wise multiplication |
+| matmul(tensor) | Matrix multiplication |
+| transpose() | Transpose tensor |
+| flatten() | Flatten to 1D array |
+| shape() | Returns [rows, cols] |
+
+Static: Tensor.zeros(r,c), Tensor.ones(r,c), Tensor.random(r,c,scale?)
+
+---
+
+## User-Friendly Utilities (fu_*)
+
+fu_tensor(data, requiresGrad?) — Create tensor from 2D array
+fu_add(a, b) — Element-wise add
+fu_mul(a, b) — Element-wise multiply
+fu_matmul(a, b) — Matrix multiply
+fu_sum(tensor) — Sum all elements
+fu_mean(tensor) — Mean of all elements
+fu_relu(tensor) — ReLU activation
+fu_sigmoid(tensor) — Sigmoid activation
+fu_tanh(tensor) — Tanh activation
+fu_softmax(tensor) — Softmax activation
+fu_flatten(tensor) — Flatten to 1D
+fu_reshape(tensor, rows, cols) — Reshape tensor
+fu_stack(tensors) — Stack tensors

package/README.md

@@ -1,4 +1,4 @@
-## Mini-JSTorch (MAJOR UPDATE)
+## Mini-JSTorch (v2.0.2)
 
 ---
 
@@ -7,15 +7,14 @@ It runs in Node.js and modern browsers, with a simple API inspired by PyTorch-st
 
 This project prioritizes `clarity`, `numerical correctness`, and `accessibility` over performance or large-scale production use.
 
-In this version `2.0.0`, we introduce:
-- **Fixed Linear layer cache** (critical bug fix for training)
-- **Fixed GELU gradient calculation**
-- **Fixed MSELoss gradient scaling**
-- **Optimized Softmax gradient** (O(n²) → O(n))
-- **Improved Tokenizer** with proper PAD/UNK separation
-- **Added Sequential.zeroGrad(), train(), eval(), stateDict() methods**
+### Changelog
 
----
+**v2.0.2:**
+- **Fixed critical training bug:** Optimizers (Adam, SGD, AdamW, Lion) now correctly update Linear and Conv2D layer weights
+- **Fixed BatchNorm2d:** Inference mode no longer produces NaN for multi-channel inputs
+- **Fixed ELU activation:** Backward pass now uses correct derivative formula
+- **Fixed saveModel/loadModel:** Now correctly saves and restores all layer types including Conv2D and BatchNorm2d
+- **Fixed BatchNorm2d gradient zeroing:** gradWeight/gradBias now correctly reset between batches
 
 **⚠️ BREAKING CHANGES in v2.0.0:**
 - Tokenizer API: `tokenizeBatch()` → `transform()`, `detokenizeBatch()` → `inverseTransform()`
@@ -80,7 +79,6 @@ In Browser/Website:
         async function train() {
             const statusEl = document.getElementById('status');
             const logEl = document.getElementById('log');
-            
             try {
                 const model = new Sequential([
                     new Linear(2, 16), new Tanh(),
@@ -107,13 +105,13 @@ In Browser/Website:
                     }
                 }
 
-                statusEl.textContent = '✅ Done';
+                statusEl.textContent = 'Done';
                 const preds = model.forward(X);
                 document.getElementById('res').innerHTML = `<h4>Results:</h4>` + 
                     X.map((input, i) => `[${input}] -> <b>${preds[i][0].toFixed(4)}</b> (Target: ${y[i][0]})`).join('<br>');
 
             } catch (e) {
-                statusEl.textContent = '❌ Error: ' + e.message;
+                statusEl.textContent = 'Error: ' + e.message;
             }
         }
         train();
@@ -202,7 +200,7 @@ git clone https://github.com/Rizal-HID11/mini-jstorch-github
 
 # Quick Start (Recommended Loss)
 
-# Multi-class Classification (SoftmaxCrossEntropy)
+## Multi-class Classification (SoftmaxCrossEntropy)
 
 ```javascript
 import {
@@ -247,7 +245,7 @@ for (let epoch = 1; epoch <= 300; epoch++) {
 ```
 `Important:` Do not combine `SoftmaxCrossEntropyLoss` with a `Softmax` layer.
 
-# Binary Classifiaction (BCEWithLogitsLoss)
+## Binary Classifiaction (BCEWithLogitsLoss)
 
 ```javascript
 import {
@@ -312,18 +310,23 @@ finalProbs.forEach((prob, i) => {
 });
 console.log(`\nAccuracy: ${(correct / X.length * 100).toFixed(2)}%`);
 ```
-Do not combine `BCEWithLogitsLoss` with a `Sigmoid` layer.
+`Important:` Do not combine `BCEWithLogitsLoss` with a `Sigmoid` layer.
 
 ---
 
 # Save & Load Models 
 
 ```javascript
-// WARN: Error/Bug may be expected for this time!
-import { saveModel, loadModel, Sequential } from "./src/jstorch.js";
+import { saveModel, loadModel, Sequential } from "./src/jstorch";
 
+// Save trained model 
 const json = saveModel(model);
-const model2 = new Sequential([...]); // same architecture
+
+// Create fresh model with same architecture and load weights 
+const model2 = new Sequential([
+    new Linear(2, 16), new ReLU(),
+    new Linear(16, 1)
+]);
 loadModel(model2, json);
 ```
 
@@ -360,7 +363,7 @@ node demo/<fileNameInDemo>.js
 
 MIT License
 
-Copyright (c) 2024
+Copyright (c) 2024-2025
 rizal-editors
 
 ---

package/demo/linear_regression.js

@@ -39,4 +39,4 @@ finalPred.forEach((p, i) => {
 });
 console.log(`\nAverage Error: ${(totalError / X.length).toFixed(2)}`);
 console.log(`Weight (slope): ${model.layers[0].W[0][0].toFixed(4)} (expected: 2.0)`);
-console.log(`Bias: ${model.layers[0].b[0].toFixed(4)} (expected: 0.0)`);
+console.log(`Bias: ${model.layers[0].b[0][0].toFixed(4)} (expected: 0.0)`);

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "mini-jstorch",
-  "version": "2.0.0",
+  "version": "2.0.2",
   "type": "module", 
   "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",
@@ -15,8 +15,7 @@
     "tiny-ml",
     "mini-neural-network",
     "mini-ml-library",
-    "mini-js-ml",
-    "educational-ml"
+    "mini-js-ml"
   ],
   "author": "Rizal",
   "license": "MIT"