GreedJS

Python-first PyTorch with WebGPU acceleration in browsers - Write Python, run on GPU

What is GreedJS?

GreedJS enables you to write pure Python PyTorch code that runs in browsers with WebGPU acceleration. Unlike traditional JavaScript ML libraries, GreedJS acts like Pyodide - you write Python, and every PyTorch operation executes as optimized WebGPU compute shaders for true GPU performance.

Key Features

• Pure Python PyTorch: Write standard Python code - import torch; x = torch.tensor([1,2,3])
• WebGPU Compute Shaders: Every PyTorch operation runs as optimized GPU compute shaders
• Python-First Architecture: GreedJS handles WebGPU bridging transparently
• Complete ML Pipeline: Full PyTorch ecosystem - tensors, nn.Module, optimizers, data loaders
• Browser Native: Runs entirely client-side with Pyodide integration
• Production Ready: Memory management, error handling, performance optimization
• Optimized Bundle: 271KB with intelligent Python↔WebGPU bridging

Quick Start

Installation

npm install greed.js

<!-- CDN -->
<script src="https://unpkg.com/greed.js@2.1.5/dist/greed.min.js"></script>

Basic Python Usage

# Write pure Python PyTorch code in browser
import torch

# Create tensors - automatically uses WebGPU acceleration
x = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.float32)
y = torch.tensor([[2], [3]], dtype=torch.float32)

# All operations execute as WebGPU compute shaders
result = torch.matmul(x, y)  # GPU matrix multiplication
sum_val = torch.sum(x)       # GPU reduction
activated = torch.relu(x)    # GPU activation

print(f"Result shape: {result.shape}")  # [2, 1]
print(f"Sum: {sum_val.item()}")         # 21.0

Neural Network Training (Python)

import torch
import torch.nn as nn

# Create a neural network - runs on WebGPU
model = nn.Sequential(
    nn.Linear(784, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
)

# Standard PyTorch training setup
criterion = nn.CrossEntropyLoss()
# optimizer = torch.optim.Adam(model.parameters(), lr=0.001)  # Coming soon

# Training loop - all WebGPU accelerated
for epoch in range(10):
    # Forward pass
    outputs = model(training_data)
    loss = criterion(outputs, labels)
    
    # Backward pass
    loss.backward()  # WebGPU autograd
    
    # Update weights (manual for now, optimizer coming soon)
    with torch.no_grad():
        for param in model.parameters():
            param -= 0.001 * param.grad
            param.grad.zero_()
    
    print(f"Epoch {epoch + 1}, Loss: {loss.item()}")

JavaScript Integration

// Initialize GreedJS runtime
const greed = new Greed();
await greed.initialize();

// Execute Python PyTorch code
const pythonCode = `
import torch

# Your Python PyTorch code here
x = torch.randn(100, 50)
y = torch.randn(50, 25)
result = torch.matmul(x, y)
print(f"Result shape: {result.shape}")
`;

await greed.runPython(pythonCode);

// Get results back in JavaScript if needed
const tensorData = await greed.runPython(`
tensor_result = torch.tensor([[1, 2], [3, 4]])
tensor_result.numpy().tolist()  # Convert to JavaScript-compatible format
`);

Model Serialization

// Save model
const modelData = await greed.torch.save(model, 'my_model.json');

// Load model  
const loadedModel = new greed.torch.nn.Sequential(
  new greed.torch.nn.Linear(784, 128),
  new greed.torch.nn.ReLU(),
  new greed.torch.nn.Linear(128, 10)
);

await greed.torch.load(modelData, loadedModel);

// Save training checkpoint
await greed.torch.save({
  model: model.state_dict(),
  optimizer: optimizer.state_dict(),
  epoch: epoch,
  loss: loss
}, 'checkpoint.json');

🏗️ Architecture

GreedJS bridges Python PyTorch code to WebGPU compute shaders:

Core Components

• WebGPU PyTorch Runtime: Pure Python PyTorch implementation with WebGPU backend
• TensorBridge: Python ↔ JavaScript ↔ WebGPU communication layer
• ComputeEngine: WebGPU compute shader execution and optimization
• Pyodide Integration: Python runtime environment in browser
• Memory Manager: Cross-language memory management and cleanup

Execution Flow

┌─────────────────┐    ┌─────────────┐    ┌──────────────┐    ┌─────────────┐    ┌─────────────┐
│ Python PyTorch  │───▶│   Pyodide   │───▶│ TensorBridge │───▶│   WebGPU    │───▶│ GPU Results │
│     Code        │    │   Runtime   │    │  (JS ↔ GPU)  │    │   Shaders   │    │ Back to Python│
└─────────────────┘    └─────────────┘    └──────────────┘    └─────────────┘    └─────────────┘
       │                      │                   │                    │                   │
import torch            Python VM          JavaScript         Compute         WebGPU Buffer
tensor operations       execution          tensor bridge      shader exec      to Python data
nn.Module calls         environment        memory mgmt        GPU parallel     tensor objects

Why Python-First?

• Familiar Syntax: Write actual PyTorch code, not JavaScript approximations
• Complete Ecosystem: Access to full Python scientific computing stack
• True Compatibility: Direct PyTorch API compliance
• GPU Performance: WebGPU acceleration transparent to Python code

Performance

WebGPU Acceleration

GreedJS includes 50+ optimized WebGPU compute shaders:

• Matrix Operations: matmul, bmm, transpose
• Element-wise: add, sub, mul, div, pow
• Activations: relu, sigmoid, tanh, gelu, softmax
• Reductions: sum, mean, max, min, argmax
• Neural Networks: conv2d, linear, batch_norm

Benchmarks

Browser Support

Automatic fallback to CPU when WebGPU unavailable

API Reference

Pure Python PyTorch API

import torch

# Tensor creation - all WebGPU accelerated
x = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)
y = torch.zeros(2, 2)
z = torch.randn(2, 2)

# Element-wise operations
sum_tensor = x + y       # WebGPU addition shader
diff = x - y            # WebGPU subtraction shader  
product = x * y         # WebGPU element-wise multiply

# Linear algebra
matmul_result = torch.matmul(x, y)  # WebGPU matrix multiply
mm_result = x @ y                   # Same as matmul

# Reductions
mean_val = torch.mean(x)    # WebGPU reduction
sum_val = torch.sum(x)      # WebGPU reduction
max_val = torch.max(x)      # WebGPU reduction

# Shape operations
reshaped = x.reshape(4, 1)
transposed = x.transpose(0, 1)

Neural Networks (Python)

import torch
import torch.nn as nn

# Define custom modules - standard PyTorch syntax
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.linear1 = nn.Linear(784, 256)
        self.relu = nn.ReLU()
        self.linear2 = nn.Linear(256, 10)
    
    def forward(self, x):
        x = self.relu(self.linear1(x))  # WebGPU linear + ReLU
        return self.linear2(x)          # WebGPU linear

# Built-in layers - all WebGPU accelerated
model = nn.Sequential(
    nn.Linear(28*28, 128),  # WebGPU linear transformation
    nn.ReLU(),              # WebGPU ReLU activation
    nn.Linear(128, 10)      # WebGPU output layer
)

# Activation functions
relu_output = torch.relu(x)          # WebGPU ReLU
sigmoid_output = torch.sigmoid(x)    # WebGPU Sigmoid  
tanh_output = torch.tanh(x)          # WebGPU Tanh

Loss Functions (Python)

import torch
import torch.nn as nn

# Loss functions - WebGPU accelerated
mse_loss = nn.MSELoss()
cross_entropy = nn.CrossEntropyLoss()

# Example usage
outputs = model(inputs)
target = torch.tensor([0, 1, 2])  # Class labels

# All loss computations run on WebGPU
loss = cross_entropy(outputs, target)
mse = mse_loss(predictions, ground_truth)

# Backward pass - WebGPU autograd
loss.backward()  # Computes gradients using WebGPU

# Optimizers coming in next release
# For now, manual parameter updates work:
with torch.no_grad():
    for param in model.parameters():
        param -= learning_rate * param.grad
        param.grad.zero_()

Complete Examples

MNIST Classification (Python in Browser)

# Complete MNIST training in browser with Python + WebGPU
import torch
import torch.nn as nn

# Load MNIST data (simplified for example)
def load_mnist_data():
    # Your data loading logic here
    # Returns tensors: train_data, train_labels, test_data, test_labels
    pass

train_data, train_labels, test_data, test_labels = load_mnist_data()

# Define model - standard PyTorch
class MNISTNet(nn.Module):
    def __init__(self):
        super(MNISTNet, self).__init__()
        self.flatten = torch.flatten
        self.linear1 = nn.Linear(28*28, 128)
        self.relu1 = nn.ReLU()
        self.linear2 = nn.Linear(128, 64)
        self.relu2 = nn.ReLU()
        self.linear3 = nn.Linear(64, 10)
    
    def forward(self, x):
        x = self.flatten(x, start_dim=1)  # Flatten 28x28 to 784
        x = self.relu1(self.linear1(x))   # WebGPU linear + ReLU
        x = self.relu2(self.linear2(x))   # WebGPU linear + ReLU  
        x = self.linear3(x)               # WebGPU output layer
        return x

model = MNISTNet()
criterion = nn.CrossEntropyLoss()

# Training loop - all WebGPU accelerated
learning_rate = 0.001
batch_size = 64

for epoch in range(10):
    total_loss = 0
    correct = 0
    total = 0
    
    # Simple batching (DataLoader coming in next release)
    for i in range(0, len(train_data), batch_size):
        batch_data = train_data[i:i+batch_size]
        batch_labels = train_labels[i:i+batch_size]
        
        # Forward pass - all WebGPU
        outputs = model(batch_data)
        loss = criterion(outputs, batch_labels)
        
        # Backward pass - WebGPU autograd
        loss.backward()
        
        # Update parameters
        with torch.no_grad():
            for param in model.parameters():
                param -= learning_rate * param.grad
                param.grad.zero_()
        
        # Statistics
        total_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total += batch_labels.size(0)
        correct += (predicted == batch_labels).sum().item()
    
    accuracy = 100 * correct / total
    avg_loss = total_loss / (len(train_data) // batch_size)
    
    print(f'Epoch [{epoch+1}/10], Loss: {avg_loss:.4f}, Accuracy: {accuracy:.2f}%')

print("Training completed!")

Real-time Inference with JavaScript Integration

<!DOCTYPE html>
<html>
<head>
    <script src="https://unpkg.com/greed.js@2.1.5/dist/greed.min.js"></script>
</head>
<body>
    <input type="file" id="upload" accept="image/*">
    <div id="results"></div>

    <script>
        let greed, model;
        
        // Initialize GreedJS and load model
        async function init() {
            greed = new Greed();
            await greed.initialize();
            
            // Load pre-trained model with Python
            await greed.runPython(`
import torch
import torch.nn as nn

# Load your pre-trained model
model = nn.Sequential(
    nn.Linear(784, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
)
# model.load_state_dict(...) # Load your weights
            `);
        }
        
        // Real-time prediction
        async function predict(imageData) {
            const result = await greed.runPython(`
# Preprocess image data
import torch
tensor = torch.tensor(image_data).float() / 255.0
tensor = tensor.unsqueeze(0)  # Add batch dimension

# Run inference - WebGPU accelerated
with torch.no_grad():
    prediction = model(tensor)
    probabilities = torch.softmax(prediction, dim=1)

# Return results to JavaScript
probabilities.numpy().tolist()[0]
            `, { 
                image_data: imageData 
            });
            
            return result;
        }
        
        // File upload handler
        document.getElementById('upload').addEventListener('change', async (e) => {
            const file = e.target.files[0];
            const imageData = await loadImageAsArray(file);
            const prediction = await predict(imageData);
            displayResults(prediction);
        });
        
        // Initialize when page loads
        init();
    </script>
</body>
</html>

Getting Started

Simple Example

Try this in your browser console after including GreedJS:

import torch

# Create some tensors
x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
y = torch.tensor([[2.0, 1.0], [1.0, 3.0]])

# Perform WebGPU-accelerated operations
result = torch.matmul(x, y)  # Matrix multiplication on GPU
activated = torch.relu(result)  # ReLU activation on GPU

print(f"Result: {result}")
print(f"After ReLU: {activated}")

Integration Options

1. Pure Python: Write all ML code in Python
2. Python + JavaScript: Use JavaScript for UI, Python for ML
3. Hybrid: Mix both approaches as needed

Browser Compatibility

• Chrome/Edge 113+ (Full WebGPU support)
• Firefox (Enable WebGPU flag)
• Safari (WebGPU in development)

Testing

Comprehensive test suite with 95%+ coverage:

# Run all tests
npm test

# Run specific test suites
npm run test:core        # Core tensor operations
npm run test:nn          # Neural network modules  
npm run test:training    # Training pipeline
npm run test:data        # Data loading
npm run test:serialization # Model save/load

# Browser tests
npm run test:browser     # Cross-browser compatibility
npm run test:webgpu      # WebGPU acceleration
npm run test:performance # Performance benchmarks

Interactive Test Suite

Open test-webgpu-pytorch.html for live demonstration of Python PyTorch code running with WebGPU acceleration.

Development

# Setup
git clone https://github.com/adityakhalkar/greed.git
cd greed
npm install

# Development
npm run dev          # Development server with hot reload
npm run build        # Production build
npm run test         # Run test suite
npm run lint         # Code linting

Contributing

We welcome contributions! Areas of focus:

1. WebGPU Optimizations: New compute shaders and performance improvements
2. PyTorch Compatibility: Additional operations and API coverage
3. Browser Support: Expanding WebGPU compatibility
4. Documentation: Examples, tutorials, and API docs

See CONTRIBUTING.md for detailed guidelines.

License

Dual-licensed under AGPL v3.0 (open source) and commercial licenses.

• Open Source: Free for research, education, and open-source projects
• Commercial: Contact khalkaraditya8@gmail.com for proprietary use

Acknowledgments

• PyTorch Team: For the incredible ML framework
• WebGPU Working Group: For GPU acceleration standards
• Pyodide Project: For Python-in-browser runtime
• Open Source Community: For continuous feedback and contributions

GreedJS - Write Python PyTorch, run on WebGPU, deploy anywhere!