EsTorch

English | 中文

PyTorch-like JavaScript (ECMAScript) tensor/autograd library with CPU training support. Zero dependencies, pure ES Modules.

Installation

npm install estorch

Quick Start

import torch from 'estorch';

// Create tensors
const a = torch.tensor([[1, 2], [3, 4]], { requiresGrad: true });
const b = torch.tensor([[5, 6], [7, 8]], { requiresGrad: true });

// Operations with autograd
const c = a.matmul(b);
const loss = c.sum();

// Backward pass
loss.backward();
console.log(a.grad); // gradients computed automatically

API Reference

Tensor Operations

| Function / Method | Description | |---|---| | torch.tensor(data, options?) | Create a tensor from nested arrays. Options: { shape, requiresGrad, device } | | torch.randn(shape, options?) | Create a tensor with standard normal random values | | torch.zeros(shape, options?) | Create a zero-filled tensor | | tensor.add(other) | Element-wise addition | | tensor.sub(other) | Element-wise subtraction | | tensor.mul(other) | Element-wise or scalar multiplication | | tensor.matmul(other) | Matrix multiplication | | tensor.tanh() | Tanh activation | | tensor.relu() | ReLU activation | | tensor.argmax(dim) | Argmax along dimension (supports dim=1 for 2D tensors) | | tensor.reshape(...shape) / tensor.view(...shape) | Reshape tensor | | tensor.flatten(startDim?) | Flatten from startDim | | tensor.clone() | Deep copy | | tensor.detach() | Detach from computation graph | | tensor.backward() | Compute gradients via backpropagation | | tensor.size(dim?) | Get shape or size of a specific dimension | | tensor.item() | Get scalar value (for 0-dim tensors) | | tensor.to(device) / tensor.cpu() | Move tensor to device | | torch.noGrad(fn) | Execute fn with gradient computation disabled |

Neural Network (torch.nn)

| Class / Function | Description | |---|---| | nn.Module | Base class for all modules. Provides parameters(), train(), eval(), stateDict(), loadStateDict(), to(device) | | nn.Parameter(data, shape) | A tensor subclass that is automatically registered as a module parameter | | nn.Linear(inFeatures, outFeatures) | Fully connected layer | | nn.Conv2d(inChannels, outChannels, kernelSize, options?) | 2D convolution layer. Options: { stride, padding } | | nn.AvgPool2d(kernelSize, options?) | 2D average pooling. Options: { stride } | | nn.CrossEntropyLoss() | Cross-entropy loss function |

Functional API (torch.nn.functional)

| Function | Description | |---|---| | functional.relu(x) | ReLU activation | | functional.tanh(x) | Tanh activation | | functional.maxPool2d(x, kernelSize, stride?) | Max pooling | | functional.avgPool2d(x, kernelSize, stride?) | Average pooling | | functional.crossEntropy(logits, labels) | Cross-entropy loss |

Optimizers (torch.optim)

| Class | Description | |---|---| | optim.Adam(params, options?) | Adam optimizer. Options: { lr, betas, eps }. Methods: zeroGrad(), step(), stateDict() | | optim.lrScheduler.StepLR(optimizer, options?) | Step learning rate scheduler. Options: { stepSize, gamma } |

Data Loading (torch.data)

| Function / Class | Description | |---|---| | data.DataLoader(dataset, options?) | Batch data loader with async iterator. Options: { batchSize, shuffle, device } | | data.randomSplit(dataset, lengths) | Randomly split a dataset into subsets | | data.mnist.MNISTDataset(imagesPath, labelsPath, options?) | MNIST dataset class. Options: { synthetic, tensorFactory, limit } | | data.mnist.prepareMNIST(root, options?) | Download/prepare MNIST data. Options: { download, synthetic } | | data.mnist.downloadMNIST(root, options?) | Download MNIST IDX files |

Utilities

| Function | Description | |---|---| | torch.device(type) | Create a device descriptor. Currently supports "cpu" only | | torch.save(obj, path) | Save object as JSON (Node.js only) | | torch.load(path) | Load JSON object from file (Node.js only) |

MNIST Training Example

See examples/mnist/main.js for a complete user example.

# Clone the repo and run with synthetic data (no download required)
git clone <your-repo-url>
cd estorch
node examples/mnist/main.js --synthetic

# Or with real MNIST data
yarn mnist:download
node examples/mnist/main.js

Example Options

node examples/mnist/main.js --synthetic --epochs 3 --batch-size 64
node examples/mnist/main.js --data-dir ./datasets/mnist --lr 0.002 --train-limit 5000

| Option | Default | Description | |--------|---------|-------------| | --synthetic | false | Use synthetic data (no download) | | --data-dir <path> | ./datasets/mnist | Path to MNIST IDX files | | --epochs <n> | 1 | Number of training epochs | | --batch-size <n> | 64 | Batch size | | --lr <n> | 0.001 | Learning rate | | --train-limit <n> | 48000 | Limit training samples | | --val-limit <n> | 12000 | Limit validation samples | | --test-limit <n> | 10000 | Limit test samples |

Smoke Tests

yarn test:smoke

Notes

• Runtime: Node.js >= 16.0.0
• Module system: ES Modules (import / export)
• Backend: CPU only (WebGPU support is experimental and not included in this package)
• Zero dependencies: No external packages required
• torch.save() and torch.load() use fs and are only available in Node.js

License

MIT