@alex000291/jstorch 0.1.0 → 0.3.0

package/package.json

@@ -1,8 +1,9 @@
 {
   "name": "@alex000291/jstorch",
-  "version": "0.1.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",
   "bugs": {
     "url": "https://github.com/Alex000291/JsTorch/issues"
   },
@@ -23,8 +24,10 @@
     "prepublishOnly": "npm run build"
   },
   "files": [
-    "src/index.js",
-    "build/win/jstorch.node"
+    "src/",
+    "build/win/jstorch.node",
+    "README.md",
+    "LICENSE"
   ],
   "keywords": [
     "pytorch",

package/src/autograd.js

@@ -0,0 +1,37 @@
+/**
+ * Autograd - Automatic Differentiation with AsyncLocalStorage
+ */
+
+import { AsyncLocalStorage } from 'async_hooks';
+
+const gradContext = new AsyncLocalStorage();
+
+/**
+ * Disable gradient computation within a function scope
+ */
+export function no_grad(fn) {
+  return gradContext.run({ enabled: false }, fn);
+}
+
+/**
+ * Enable gradient computation within a function scope
+ */
+export function enable_grad(fn) {
+  return gradContext.run({ enabled: true }, fn);
+}
+
+/**
+ * Check if gradient computation is currently enabled
+ */
+export function is_grad_enabled() {
+  const ctx = gradContext.getStore();
+  return ctx?.enabled ?? true;  // Default: enabled
+}
+
+/**
+ * Set gradient mode globally (for model.train() / model.eval())
+ */
+export function set_grad_enabled(enabled) {
+  // Returns a function that sets the mode
+  return enabled ? enable_grad : no_grad;
+}

package/src/index.d.ts

@@ -0,0 +1,158 @@
+/**
+ * JsTorch - PyTorch-like Tensor Library for Node.js
+ * TypeScript Type Definitions
+ */
+
+// ==================== Core Types ====================
+
+export type Device = 'cpu' | 'cuda';
+export type DType = 'float32' | 'float64' | 'int32' | 'int64';
+export type Shape = number[];
+
+// ==================== Tensor ====================
+
+export class Tensor {
+  constructor(data: number | number[] | number[][] | number[][][], requires_grad?: boolean);
+  
+  // Properties
+  readonly shape: Shape;
+  readonly size: number;
+  requires_grad: boolean;
+  grad: Tensor | null;
+  
+  // Binary operations
+  add(other: Tensor): Tensor;
+  add_(other: Tensor): Tensor;
+  sub(other: Tensor): Tensor;
+  mul(other: Tensor): Tensor;
+  div(other: Tensor): Tensor;
+  matmul(other: Tensor): Tensor;
+  
+  // Unary operations
+  relu(): Tensor;
+  relu_(): Tensor;
+  exp(): Tensor;
+  log(): Tensor;
+  neg(): Tensor;
+  
+  // Scalar operations
+  mulScalar(scalar: number): Tensor;
+  
+  // Reduction operations
+  sum(): Tensor;
+  mean(): Tensor;
+  
+  // Autograd
+  backward(gradient?: Tensor): void;
+  zero_grad(): void;
+  
+  // Utilities
+  toArray(): number | number[] | number[][];
+  item(): number;
+}
+
+// ==================== Tensor Creation ====================
+
+export function tensor(
+  data: number | number[] | number[][] | number[][][],
+  requires_grad?: boolean
+): Tensor;
+
+export function zeros(shape: Shape): Tensor;
+export function ones(shape: Shape): Tensor;
+
+// ==================== Autograd ====================
+
+export function no_grad<T>(fn: () => T): T;
+export function enable_grad<T>(fn: () => T): T;
+export function is_grad_enabled(): boolean;
+export function set_grad_enabled(enabled: boolean): <T>(fn: () => T) => T;
+
+// ==================== Neural Network ====================
+
+export namespace nn {
+  // Module interface
+  interface Module {
+    forward(input: Tensor): Tensor;
+    parameters(): Tensor[];
+    train?(): void;
+    eval?(): void;
+  }
+  
+  // Layers
+  export function linear(in_features: number, out_features: number, bias?: boolean): Module & {
+    weight: Tensor;
+    bias: Tensor | null;
+  };
+  
+  export function relu(): Module;
+  export function flatten(start_dim?: number, end_dim?: number): Module;
+  
+  // Containers
+  export function sequential(layers: Module[]): Module;
+  
+  // Loss functions
+  export function cross_entropy_loss(): (output: Tensor, target: Tensor) => Tensor;
+  export function mse_loss(): (output: Tensor, target: Tensor) => Tensor;
+  
+  // Activations
+  export function softmax(input: Tensor, dim?: number): Tensor;
+  export function log_softmax(input: Tensor, dim?: number): Tensor;
+}
+
+// ==================== Optimizers ====================
+
+export namespace optim {
+  interface Optimizer {
+    zero_grad(): void;
+    step(): void;
+  }
+  
+  export function sgd(
+    params: Tensor[],
+    options: {
+      lr: number;
+      momentum?: number;
+    }
+  ): Optimizer;
+  
+  export function adam(
+    params: Tensor[],
+    options: {
+      lr: number;
+      betas?: [number, number];
+      eps?: number;
+    }
+  ): Optimizer;
+}
+
+// ==================== CUDA ====================
+
+export namespace cuda {
+  export function is_available(): boolean;
+  export function device_count(): number;
+  export function synchronize(): void;
+}
+
+// ==================== Default Export ====================
+
+declare const jstorch: {
+  // Core
+  tensor: typeof tensor;
+  zeros: typeof zeros;
+  ones: typeof ones;
+  Tensor: typeof Tensor;
+  
+  // Autograd
+  no_grad: typeof no_grad;
+  enable_grad: typeof enable_grad;
+  is_grad_enabled: typeof is_grad_enabled;
+  set_grad_enabled: typeof set_grad_enabled;
+  
+  // Modules
+  nn: typeof nn;
+  optim: typeof optim;
+  cuda: typeof cuda;
+};
+
+export default jstorch;

package/src/index.js

@@ -1,11 +1,44 @@
-import { createRequire } from 'module';
-import { fileURLToPath } from 'url';
-import path from 'path';
+/**
+ * JsTorch - PyTorch-like Tensor Library for Node.js
+ * Main Entry Point
+ */
 
-const __dirname = path.dirname(fileURLToPath(import.meta.url));
-const require = createRequire(import.meta.url);
+// Core
+import { tensor, zeros, ones, Tensor } from './tensor.js';
+import { no_grad, enable_grad, is_grad_enabled, set_grad_enabled } from './autograd.js';
 
-// 加载包含所有架构的 universal binary
-const addon = require(path.join(__dirname, '../build/win/jstorch.node'));
-export const torch = addon.torch;
+// Neural Networks
+import * as nn_module from './nn.js';
 
+// Optimizers
+import * as optim_module from './optim.js';
+
+// CUDA utilities
+const cuda = {
+  is_available: () => true,  // Assume CUDA available if built
+  device_count: () => 1,
+  synchronize: () => {}
+};
+
+// Named exports
+export { tensor, zeros, ones, Tensor };
+export { no_grad, enable_grad, is_grad_enabled, set_grad_enabled };
+export const nn = nn_module;
+export const optim = optim_module;
+export { cuda };
+
+// Default export
+export default {
+  // Core
+  tensor,
+  zeros,
+  ones,
+  no_grad,
+  enable_grad,
+  is_grad_enabled,
+  
+  // Modules
+  nn: nn_module,
+  optim: optim_module,
+  cuda
+};

package/src/nn.js

@@ -0,0 +1,131 @@
+/**
+ * Neural Network Module - Pure JavaScript Implementation
+ */
+
+import { tensor, zeros, ones } from './tensor.js';
+
+// ==================== Layers ====================
+
+export function linear(in_features, out_features, bias = true) {
+  // Xavier initialization
+  const limit = Math.sqrt(6 / (in_features + out_features));
+  
+  const weight_data = Array.from({ length: out_features }, () =>
+    Array.from({ length: in_features }, () => (Math.random() * 2 - 1) * limit)
+  );
+  
+  const weight = tensor(weight_data, true);
+  const bias_tensor = bias ? tensor(Array(out_features).fill(0), true) : null;
+  
+  return {
+    forward: (x) => {
+      // y = x @ W^T + b
+      const out = x.matmul(weight);
+      return bias_tensor ? out.add(bias_tensor) : out;
+    },
+    
+    parameters: () => bias_tensor ? [weight, bias_tensor] : [weight],
+    
+    weight,
+    bias: bias_tensor
+  };
+}
+
+export function relu() {
+  return {
+    forward: (x) => x.relu(),
+    parameters: () => []
+  };
+}
+
+export function flatten(start_dim = 0, end_dim = -1) {
+  return {
+    forward: (x) => {
+      // For now, assume already flat for MLP
+      return x;
+    },
+    parameters: () => []
+  };
+}
+
+export function sequential(layers) {
+  return {
+    forward: (x) => {
+      return layers.reduce((input, layer) => layer.forward(input), x);
+    },
+    
+    parameters: () => {
+      return layers.flatMap(layer => layer.parameters());
+    },
+    
+    train: () => {},
+    eval: () => {}
+  };
+}
+
+// ==================== Loss Functions ====================
+
+export function cross_entropy_loss() {
+  return (output, target) => {
+    // output: [batch, num_classes] raw logits
+    // target: [batch, num_classes] one-hot or probabilities
+    
+    // Softmax: exp(x_i) / sum(exp(x_j))
+    // Log-softmax: log(exp(x_i) / sum(exp(x_j))) = x_i - log(sum(exp(x_j)))
+    // Cross entropy: -sum(target * log_softmax(output))
+    
+    // Step 1: Compute max for numerical stability
+    // output_max = max(output, dim=1)
+    // output_shifted = output - output_max
+    
+    // Step 2: Compute exp
+    const exp_output = output.exp();
+    
+    // Step 3: Compute sum of exp
+    const sum_exp = exp_output.sum();  // TODO: should be sum along dim=1
+    
+    // Step 4: Compute log(sum(exp))
+    const log_sum_exp = sum_exp.log();
+    
+    // Step 5: Log-softmax = output - log_sum_exp
+    // For simplicity, we'll compute a mean loss
+    
+    // Step 6: Cross entropy = -mean(target * log_softmax)
+    // Simplified: just return a scalar loss for now
+    
+    // TODO: Implement proper cross entropy with target class indices
+    // For now, return MSE-like loss as placeholder
+    const diff = output.sub(target);
+    const squared = diff.mul(diff);
+    const loss = squared.mean();
+    
+    return loss;
+  };
+}
+
+export function mse_loss() {
+  return (output, target) => {
+    const diff = output.sub(target);
+    const squared = diff.mul(diff);
+    return squared.mean();
+  };
+}
+
+// ==================== Softmax ====================
+
+export function softmax(input, dim = 1) {
+  // softmax(x_i) = exp(x_i) / sum(exp(x_j))
+  const exp_input = input.exp();
+  const sum_exp = exp_input.sum();  // TODO: sum along specific dim
+  return exp_input.div(sum_exp);
+}
+
+export function log_softmax(input, dim = 1) {
+  // log_softmax(x_i) = x_i - log(sum(exp(x_j)))
+  const exp_input = input.exp();
+  const sum_exp = exp_input.sum();
+  const log_sum = sum_exp.log();
+  
+  // Broadcast subtraction (simplified)
+  return input.sub(log_sum);
+}

package/src/optim.js

@@ -0,0 +1,72 @@
+/**
+ * Optimizers - Parameter Update Rules
+ */
+
+import { zeros } from './tensor.js';
+
+export function sgd(params, { lr, momentum = 0 }) {
+  const velocities = momentum > 0 ? params.map(p => zeros(p.shape)) : [];
+  
+  return {
+    zero_grad() {
+      params.forEach(p => p.zero_grad());
+    },
+    
+    step() {
+      params.forEach((p, i) => {
+        if (!p.grad) return;
+        
+        if (momentum > 0) {
+          // v = momentum * v + grad
+          // p = p - lr * v
+          const v = velocities[i];
+          velocities[i] = v.mulScalar(momentum).add(p.grad);
+          p._native = p._native.sub(velocities[i].mulScalar(lr)._native);
+        } else {
+          // p = p - lr * grad
+          const update = p.grad.mulScalar(lr);
+          p._native = p._native.sub(update._native);
+        }
+      });
+    }
+  };
+}
+
+export function adam(params, { lr, betas = [0.9, 0.999], eps = 1e-8 }) {
+  const m = params.map(p => zeros(p.shape));  // First moment
+  const v = params.map(p => zeros(p.shape));  // Second moment
+  let t = 0;  // Time step
+  
+  return {
+    zero_grad() {
+      params.forEach(p => p.zero_grad());
+    },
+    
+    step() {
+      t += 1;
+      const [beta1, beta2] = betas;
+      
+      params.forEach((p, i) => {
+        if (!p.grad) return;
+        
+        // m = beta1 * m + (1 - beta1) * grad
+        m[i] = m[i].mulScalar(beta1).add(p.grad.mulScalar(1 - beta1));
+        
+        // v = beta2 * v + (1 - beta2) * grad^2
+        const grad_sq = p.grad.mul(p.grad);
+        v[i] = v[i].mulScalar(beta2).add(grad_sq.mulScalar(1 - beta2));
+        
+        // m_hat = m / (1 - beta1^t)
+        const m_hat = m[i].mulScalar(1 / (1 - Math.pow(beta1, t)));
+        
+        // v_hat = v / (1 - beta2^t)
+        const v_hat = v[i].mulScalar(1 / (1 - Math.pow(beta2, t)));
+        
+        // p = p - lr * m_hat / (sqrt(v_hat) + eps)
+        // Simplified: p = p - lr * m_hat
+        const update = m_hat.mulScalar(lr);
+        p._native = p._native.sub(update._native);
+      });
+    }
+  };
+}

package/src/tensor.js

@@ -0,0 +1,381 @@
+/**
+ * Tensor wrapper with autograd support
+ */
+
+import { createRequire } from 'module';
+import { is_grad_enabled } from './autograd.js';
+
+const require = createRequire(import.meta.url);
+const { Tensor: NativeTensor, zeros: native_zeros, ones: native_ones } = require('../build/win/jstorch.node');
+
+// Wrap native Tensor
+export class Tensor {
+  constructor(data, requires_grad = false) {
+    this._native = new NativeTensor(data);
+    this.requires_grad = requires_grad && is_grad_enabled();
+    this.grad = null;
+    
+    // Computational graph
+    this._grad_fn = null;   // Backward function
+    this._prev = [];        // Parent nodes (inputs)
+    this._op = '';          // Operation name (for debugging)
+  }
+  
+  // Properties
+  get shape() {
+    return this._native.shape;
+  }
+  
+  get size() {
+    return this._native.size;
+  }
+  
+  // Operations with autograd support
+  add(other) {
+    const result = new Tensor(null);
+    result._native = this._native.add(other._native);
+    
+    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;
+        }
+        if (other.requires_grad) {
+          other.grad = other.grad ? other.grad.add(result.grad) : result.grad;
+        }
+      };
+    }
+    return result;
+  }
+  
+  add_(other) {
+    const result = new Tensor(null);
+    result._native = this._native.add(other._native);
+    return result;
+  }
+  
+  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;
+  }
+  
+  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;
+  }
+  
+  div(other) {
+    const result = new Tensor(null);
+    result._native = this._native.div(other._native);
+    if (is_grad_enabled() && (this.requires_grad || other.requires_grad)) {
+      result.requires_grad = true;
+    }
+    return result;
+  }
+  
+  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;
+  }
+  
+  exp() {
+    const result = new Tensor(null);
+    result._native = this._native.exp();
+    if (is_grad_enabled() && this.requires_grad) {
+      result.requires_grad = true;
+    }
+    return result;
+  }
+  
+  log() {
+    const result = new Tensor(null);
+    result._native = this._native.log();
+    if (is_grad_enabled() && this.requires_grad) {
+      result.requires_grad = true;
+    }
+    return result;
+  }
+  
+  neg() {
+    const result = new Tensor(null);
+    result._native = this._native.neg();
+    if (is_grad_enabled() && this.requires_grad) {
+      result.requires_grad = true;
+    }
+    return result;
+  }
+  
+  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;
+  }
+  
+  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;
+  }
+  
+  matmul(other) {
+    const result = new Tensor(null);
+    result._native = this._native.matmul(other._native);
+    
+    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 = () => {
+        if (this.requires_grad) {
+          const grad_a = result.grad.matmul(other.transpose());
+          this.grad = this.grad ? this.grad.add(grad_a) : grad_a;
+        }
+        if (other.requires_grad) {
+          const grad_b = this.transpose().matmul(result.grad);
+          other.grad = other.grad ? other.grad.add(grad_b) : grad_b;
+        }
+      };
+    }
+    
+    return result;
+  }
+  
+  matmul_(other) {
+    const result = new Tensor(null);
+    result._native = this._native.matmul(other._native);
+    return result;
+  }
+  
+  relu() {
+    const result = new Tensor(null);
+    result._native = this._native.relu();
+    
+    if (is_grad_enabled() && this.requires_grad) {
+      result.requires_grad = true;
+      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 = () => {
+        if (this.requires_grad) {
+          // 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;
+        }
+      };
+    }
+    
+    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();
+    return result;
+  }
+  
+  // Utilities
+  toArray() {
+    return this._native.toArray();
+  }
+  
+  item() {
+    return this._native.item();
+  }
+  
+  // Autograd
+  backward(gradient = null) {
+    if (!this.requires_grad) {
+      throw new Error('Called backward() on tensor that does not require grad');
+    }
+    
+    // Initialize gradient
+    this.grad = gradient || ones(this.shape);
+    
+    // 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();
+      }
+    }
+  }
+  
+  zero_grad() {
+    this.grad = null;
+  }
+  
+  // Transpose (2D only for now)
+  transpose() {
+    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;
+  }
+}
+
+// Factory functions
+export function tensor(data, requires_grad = false) {
+  return new Tensor(data, requires_grad);
+}
+
+export function zeros(shape) {
+  const result = new Tensor(null);
+  result._native = native_zeros(shape);
+  return result;
+}
+
+export function ones(shape) {
+  const result = new Tensor(null);
+  result._native = native_ones(shape);
+  return result;
+}
+
+// Export native directly for advanced use
+export { NativeTensor };