@alex000291/jstorch 0.2.0 → 0.3.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@alex000291/jstorch",
-  "version": "0.2.0",
+  "version": "0.3.0",
   "description": "A JavaScript version of PyTorch with CUDA acceleration",
   "homepage": "https://github.com/Alex000291/JsTorch#readme",
   "types": "src/index.d.ts",

package/src/tensor.js

@@ -14,7 +14,11 @@ export class Tensor {
     this._native = new NativeTensor(data);
     this.requires_grad = requires_grad && is_grad_enabled();
     this.grad = null;
-    this._grad_fn = null;
+    
+    // Computational graph
+    this._grad_fn = null;   // Backward function
+    this._prev = [];        // Parent nodes (inputs)
+    this._op = '';          // Operation name (for debugging)
   }
   
   // Properties
@@ -33,6 +37,10 @@ export class Tensor {
     
     if (is_grad_enabled() && (this.requires_grad || other.requires_grad)) {
       result.requires_grad = true;
+      result._prev = [this, other];
+      result._op = 'add';
+      
+      // Gradient function: d(a+b)/da = 1, d(a+b)/db = 1
       result._grad_fn = () => {
         if (this.requires_grad) {
           this.grad = this.grad ? this.grad.add(result.grad) : result.grad;
@@ -54,8 +62,22 @@ export class Tensor {
   sub(other) {
     const result = new Tensor(null);
     result._native = this._native.sub(other._native);
+    
     if (is_grad_enabled() && (this.requires_grad || other.requires_grad)) {
       result.requires_grad = true;
+      result._prev = [this, other];
+      result._op = 'sub';
+      
+      // d(a-b)/da = 1, d(a-b)/db = -1
+      result._grad_fn = () => {
+        if (this.requires_grad) {
+          this.grad = this.grad ? this.grad.add(result.grad) : result.grad;
+        }
+        if (other.requires_grad) {
+          const neg_grad = result.grad.neg();
+          other.grad = other.grad ? other.grad.add(neg_grad) : neg_grad;
+        }
+      };
     }
     return result;
   }
@@ -63,8 +85,23 @@ export class Tensor {
   mul(other) {
     const result = new Tensor(null);
     result._native = this._native.mul(other._native);
+    
     if (is_grad_enabled() && (this.requires_grad || other.requires_grad)) {
       result.requires_grad = true;
+      result._prev = [this, other];
+      result._op = 'mul';
+      
+      // d(a*b)/da = b, d(a*b)/db = a
+      result._grad_fn = () => {
+        if (this.requires_grad) {
+          const grad_a = result.grad.mul(other);
+          this.grad = this.grad ? this.grad.add(grad_a) : grad_a;
+        }
+        if (other.requires_grad) {
+          const grad_b = result.grad.mul(this);
+          other.grad = other.grad ? other.grad.add(grad_b) : grad_b;
+        }
+      };
     }
     return result;
   }
@@ -81,8 +118,19 @@ export class Tensor {
   mulScalar(scalar) {
     const result = new Tensor(null);
     result._native = this._native.mulScalar(scalar);
+    
     if (is_grad_enabled() && this.requires_grad) {
       result.requires_grad = true;
+      result._prev = [this];
+      result._op = 'mulScalar';
+      
+      // d(a*c)/da = c (where c is scalar)
+      result._grad_fn = () => {
+        if (this.requires_grad) {
+          const grad_a = result.grad.mulScalar(scalar);
+          this.grad = this.grad ? this.grad.add(grad_a) : grad_a;
+        }
+      };
     }
     return result;
   }
@@ -117,8 +165,21 @@ export class Tensor {
   sum() {
     const result = new Tensor(null);
     result._native = this._native.sum();
+    
     if (is_grad_enabled() && this.requires_grad) {
       result.requires_grad = true;
+      result._prev = [this];
+      result._op = 'sum';
+      
+      // d(sum(a))/da = ones_like(a) * grad_output
+      result._grad_fn = () => {
+        if (this.requires_grad) {
+          // Gradient broadcasts to all elements
+          const grad_value = result.grad.item();
+          const grad_a = ones(this.shape).mulScalar(grad_value);
+          this.grad = this.grad ? this.grad.add(grad_a) : grad_a;
+        }
+      };
     }
     return result;
   }
@@ -126,8 +187,20 @@ export class Tensor {
   mean() {
     const result = new Tensor(null);
     result._native = this._native.mean();
+    
     if (is_grad_enabled() && this.requires_grad) {
       result.requires_grad = true;
+      result._prev = [this];
+      result._op = 'mean';
+      
+      // d(mean(a))/da = ones_like(a) / size * grad_output
+      result._grad_fn = () => {
+        if (this.requires_grad) {
+          const grad_value = result.grad.item();
+          const grad_a = ones(this.shape).mulScalar(grad_value / this.size);
+          this.grad = this.grad ? this.grad.add(grad_a) : grad_a;
+        }
+      };
     }
     return result;
   }
@@ -138,8 +211,11 @@ export class Tensor {
     
     if (is_grad_enabled() && (this.requires_grad || other.requires_grad)) {
       result.requires_grad = true;
+      result._prev = [this, other];
+      result._op = 'matmul';
+      
+      // d(A@B)/dA = dC @ B^T, d(A@B)/dB = A^T @ dC
       result._grad_fn = () => {
-        // Backward: dL/dA = dL/dC @ B^T, dL/dB = A^T @ dL/dC
         if (this.requires_grad) {
           const grad_a = result.grad.matmul(other.transpose());
           this.grad = this.grad ? this.grad.add(grad_a) : grad_a;
@@ -166,12 +242,19 @@ export class Tensor {
     
     if (is_grad_enabled() && this.requires_grad) {
       result.requires_grad = true;
-      const input_data = this.toArray();  // Save for backward
+      result._prev = [this];
+      result._op = 'relu';
+      
+      // Save input for backward
+      const input_data = this.toArray();
+      
+      // d(relu(a))/da = (a > 0) ? 1 : 0
       result._grad_fn = () => {
-        // Backward: gradient passes through where input > 0
-        // TODO: implement element-wise multiplication with mask
         if (this.requires_grad) {
-          this.grad = result.grad;  // Simplified for now
+          // Create mask: 1 where input > 0, 0 otherwise
+          const mask = this._createMask(input_data, x => x > 0 ? 1 : 0);
+          const grad_a = result.grad.mul(mask);
+          this.grad = this.grad ? this.grad.add(grad_a) : grad_a;
         }
       };
     }
@@ -179,6 +262,17 @@ export class Tensor {
     return result;
   }
   
+  // Helper: create mask tensor from nested array
+  _createMask(data, fn) {
+    const applyMask = (arr) => {
+      if (Array.isArray(arr)) {
+        return arr.map(applyMask);
+      }
+      return fn(arr);
+    };
+    return tensor(applyMask(data));
+  }
+  
   relu_() {
     const result = new Tensor(null);
     result._native = this._native.relu();
@@ -203,9 +297,29 @@ export class Tensor {
     // Initialize gradient
     this.grad = gradient || ones(this.shape);
     
-    // Topological sort and execute grad_fn (simplified - just call directly for now)
-    if (this._grad_fn) {
-      this._grad_fn();
+    // Build topological order
+    const topo = [];
+    const visited = new Set();
+    
+    const build_topo = (node) => {
+      if (!node.requires_grad || visited.has(node)) return;
+      visited.add(node);
+      
+      for (const parent of node._prev) {
+        build_topo(parent);
+      }
+      
+      topo.push(node);
+    };
+    
+    build_topo(this);
+    
+    // Execute gradients in reverse topological order
+    for (let i = topo.length - 1; i >= 0; i--) {
+      const node = topo[i];
+      if (node._grad_fn) {
+        node._grad_fn();
+      }
     }
   }
   
@@ -213,10 +327,36 @@ export class Tensor {
     this.grad = null;
   }
   
-  // Placeholder for transpose (TODO: implement in C++)
+  // Transpose (2D only for now)
   transpose() {
-    // Assume 2D for now
-    return this;  // TODO
+    if (this.shape.length !== 2) {
+      throw new Error('transpose() only supports 2D tensors for now');
+    }
+    
+    const data = this.toArray();
+    const [rows, cols] = this.shape;
+    
+    // Transpose data
+    const transposed = Array.from({ length: cols }, (_, i) =>
+      Array.from({ length: rows }, (_, j) => data[j][i])
+    );
+    
+    const result = tensor(transposed, this.requires_grad && is_grad_enabled());
+    
+    if (is_grad_enabled() && this.requires_grad) {
+      result._prev = [this];
+      result._op = 'transpose';
+      
+      // d(A^T)/dA = (dC)^T
+      result._grad_fn = () => {
+        if (this.requires_grad) {
+          const grad_a = result.grad.transpose();
+          this.grad = this.grad ? this.grad.add(grad_a) : grad_a;
+        }
+      };
+    }
+    
+    return result;
   }
 }