framework

Open source

A machine learning library with a TypeScript API and Rust backend. CUDA and WebGPU compatibility. Built to understand how ML frameworks and models work internal

974 stars0 likes1 views

Website GitHub

About

A machine learning library with a TypeScript API and Rust backend. CUDA and WebGPU compatibility. Built to understand how ML frameworks and models work internal

@mni-ml/framework

A TypeScript ML framework with Rust native backends (CPU, CUDA, WebGPU) providing autograd, tensor operations, and neural network training at GPU speed.

Features

Automatic differentiation -- full backward pass through an autograd tape
GPU acceleration -- CUDA (NVIDIA) and WebGPU (Metal/Vulkan/DX12) backends
PyTorch-like API -- familiar Tensor, Module, Parameter, optimizer classes
Comprehensive ops -- elementwise, matmul, conv1d/conv2d, pooling, reductions, activations
Built-in modules -- Linear, Conv1d, Conv2d, Embedding, ReLU, Sigmoid, Tanh
Optimizers -- SGD and Adam (AdamW) with learning rate scheduling

Installation

bash

npm install @mni-ml/framework

Quick Start

…

API Reference

Tensor

…

Neural Network Modules

typescript

import { Linear, Conv1d, Conv2d, ReLU, Sigmoid, Embedding } from '@mni-ml/framework';

const linear = new Linear(inputSize, outputSize);
const conv1d = new Conv1d(inChannels, outChannels, kernelSize, stride, padding);
const conv2d = new Conv2d(inChannels, outChannels, kernelSize, stride, padding);
const embedding = new Embedding(vocabSize, embeddingDim);

// Use in forward pass
const out = linear.forward(input);

Functional Operations

typescript

import { softmax, gelu, layerNorm, crossEntropyLoss, dropout,
         avgpool2d, maxpool2d, tile } from '@mni-ml/framework';

const sm = softmax(logits, dim);
const g = gelu(x);
const ln = layerNorm(x, gamma, beta, eps);
const loss = crossEntropyLoss(logits, targets);
const dropped = dropout(x, rate, training);
const pooled = avgpool2d(x, kernelH, kernelW);
const maxPooled = maxpool2d(x, kernelH, kernelW);
const tiled = tile(x, [2, 1]);

Optimizers

typescript

import { Adam, SGD } from '@mni-ml/framework';

const optimizer = new Adam(parameters, lr, beta1, beta2, eps, weightDecay);
// or
const optimizer = new SGD(parameters, lr);

optimizer.step();      // update parameters
optimizer.zeroGrad();  // clear gradients

Backend Architecture

TypeScript API (tensor.ts, nn.ts, optimizer.ts)
    │
    └─→ N-API Bridge (lib.rs)
            │
            ├─→ CPU Backend (Vec, pure Rust)
            ├─→ CUDA Backend (cudarc + .cu kernels)
            └─→ WebGPU Backend (wgpu + .wgsl shaders)

All three backends share the same autograd tape and tensor store. Feature flags are mutually exclusive at compile time:

cpu -- default, no GPU required
cuda -- NVIDIA GPU via CUDA
webgpu -- any GPU via wgpu (Metal, Vulkan, DX12)

Building from source

Only needed if you are contributing or want a custom build. Requires Rust.

bash

# CPU (default)
npm run build:native

# CUDA (requires CUDA toolkit)
npm run build:native:cuda

# WebGPU
npm run build:native:webgpu

# Build TypeScript
npm run build