A lightweight WebGPU-powered tensor runtime for running AI workloads directly in the browser.

No server. No Python. No CUDA. Pure GPU math from JavaScript.

import { tensor } from 'tensorlite'

const a = tensor([1, 2, 3])
const b = tensor([4, 5, 6])

// Operations are lazy — no GPU work yet
const c = a.add(b).relu()

// .data() triggers GPU execution and reads back the result
console.log(await c.data())  // Float32Array [5, 7, 9]

// Matrix multiplication
const A = tensor([[1, 2], [3, 4]])          // 2×2
const B = tensor([[2, 0], [1, 3]])          // 2×2
const C = await A.matmul(B).data()

// Softmax over logits
const probs = await tensor([-1, 0, 1, 2]).softmax().data()
// probs sums to 1.0

// Chain any number of ops — one GPU round-trip per .data() call
const result = await a.mul(b).add(tensor([1,1,1])).relu().softmax().data()

# Via npm (once published)
npm install tensorlite

# Or clone and serve directly — no build step needed
git clone https://github.com/your-org/tensorlite.js
cd tensorlite.js
npm install
npm run dev

Open http://localhost:5173/examples/vector-add.html in a WebGPU browser.

Create a tensor from a JS array, TypedArray, or scalar.

tensor([1, 2, 3])                          // shape [3], dtype f32
tensor([[1, 0], [0, 1]])                   // shape [2, 2]
tensor([1,2,3,4], { shape: [2, 2] })       // explicit shape
tensor([1, 2, 3], { dtype: 'i32' })        // integer tensor

t.shape    // number[]  — e.g. [3, 4]
t.dtype    // string    — 'f32' | 'i32' | 'u32'
t.ndim     // number    — number of dimensions
t.size     // number    — total element count

a.add(b)        // element-wise addition
a.sub(b)        // element-wise subtraction
a.mul(b)        // element-wise multiplication
a.matmul(b)     // matrix multiplication  A @ B
a.relu()        // max(0, x) activation
a.softmax()     // row-wise softmax over last axis
a.sum()         // sum all elements → shape [1]
a.sum(0)        // reduce along axis 0
a.sum(-1)       // reduce along last axis

await t.data()      // TypedArray (Float32Array / Int32Array / Uint32Array)
await t.toArray()   // plain JS number[]

tensorlite.js/
├── src/
│   ├── index.js              ← public API
│   ├── core/
│   │   ├── tensor.js         ← Tensor class + tensor() factory
│   │   ├── buffer.js         ← GPU buffer upload / readback
│   │   ├── kernel.js         ← compiled WGSL kernel wrapper
│   │   └── scheduler.js      ← command batching (Phase 2)
│   ├── ops/
│   │   ├── add.js            ← element-wise add (WGSL + execute)
│   │   ├── mul.js            ← element-wise multiply
│   │   ├── matmul.js         ← matrix multiplication
│   │   ├── relu.js           ← ReLU activation
│   │   ├── softmax.js        ← row-wise softmax
│   │   └── sum.js            ← reduction sum
│   ├── webgpu/
│   │   ├── adapter.js        ← singleton GPUDevice
│   │   ├── pipeline.js       ← pipeline + bind group cache
│   │   └── shaderCompiler.js ← WGSL module cache
│   └── utils/
│       ├── dtype.js          ← DType enum + TypedArray helpers
│       ├── shape.js          ← numel, broadcast, strides…
│       └── memory.js         ← MemoryPool (Phase 2)
├── examples/
│   ├── vector-add.html
│   ├── matrix-multiply.html
│   └── image-filter.html
├── docs/
│   ├── index.html
│   ├── concepts.html
│   └── api.html
└── tests/
    └── tensor.test.js

• Tensor creation, shapes, dtypes
• Element-wise ops: add, sub, mul, relu
• Matrix multiplication (naive WGSL)
• Softmax (numerically stable, row-wise)
• Sum reduction (axis or global)
• GPU buffer upload / readback
• Lazy evaluation with deferred op graph
• Shader and pipeline caching
• Browser examples
• Full test suite

• Full NumPy-style broadcasting in WGSL (not just modular index)
• GPU memory pool — buffer reuse across ops
• Tiled matmul with workgroup shared memory (10-50× speedup)
• Native conv2d kernel
• Tensor slicing and indexing
• Transpose
• Command batching in Scheduler (reduces CPU→GPU submission overhead)
• f16 support (via WebGPU shader-f16 extension)

• Embedding lookup (gather)
• Scaled dot-product attention
• Layer normalisation
• Linear layers (forward + backward)
• Basic ONNX model loading
• Node.js / Dawn support

Operations build a computation graph without touching the GPU:

tensor([1,2,3])            → Tensor { cpuData: Float32Array }
  .add(tensor([4,5,6]))    → Tensor { pendingOp: { type:'add', inputs:[a,b] } }
  .relu()                  → Tensor { pendingOp: { type:'relu', inputs:[c] } }

When .data() is called, _realise() recurses through the graph:

1. Upload leaf tensor data to GPU storage buffers
2. Execute each pending kernel in dependency order
3. Read the final output buffer back to the CPU

This means one CPU↔GPU round-trip no matter how long the chain.

• First call is slow — WebGPU adapter initialisation + shader compilation takes ~100–500ms. All subsequent calls reuse cached pipelines.
• Workgroup sizes — ops use 64 threads/group (1D) or 8×8 (matmul). Phase 2 will tune these based on GPU limits.
• Memory transfers — every .data() call copies data CPU↔GPU. Keep tensors on the GPU as long as possible; only call .data() at the end.

Each GPU operation lives in its own file under src/ops/. Adding a new op:

1. Create src/ops/myop.js with a WGSL shader string and executeMyop() function
2. Add the 'myop' case to _dispatch() in src/core/tensor.js
3. Expose t.myop() as a method on the Tensor class
4. Add tests in tests/tensor.test.js

MIT © tensorlite contributors