@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

npm install @mni-ml/framework

Quick Start

import{Tensor,Linear,Adam,Parameter,softmax,crossEntropyLoss}from'@mni-ml/framework';// Create tensorsconstx=Tensor.rand([32,10]);consttargets=[[0],[1],[2],/* ... */];// Build a modelconstlayer1=newLinear(10,64);constlayer2=newLinear(64,3);// Forward passleth=layer1.forward(x).relu();letlogits=layer2.forward(h);letloss=crossEntropyLoss(logits,targets);// Backward passloss.backward();// Optimizeconstparams=[...layer1.parameters(), ...layer2.parameters()];constoptimizer=newAdam(params,0.001);optimizer.step();optimizer.zeroGrad();

API Reference

Tensor

// CreationTensor.zeros([2,3])// zero-filledTensor.ones([2,3])// one-filledTensor.rand([2,3])// uniform [0, 1)Tensor.randn([2,3])// normal distributionTensor.fromFloat32(data,shape)// from Float32Array// Arithmetic (with autograd)a.add(b)a.add(2.0)// additiona.sub(b)// subtractiona.mul(b)a.mul(2.0)// multiplicationa.div(b)a.div(2.0)// divisiona.neg()// negationa.exp()a.log()// exponentialsa.pow(2)// power// Activationsa.relu()a.sigmoid()// Reductionsa.sum(dim)a.sum()// sum along dim or alla.mean(dim)a.mean()// mean along dim or alla.max(dim)// max along dim// Comparisons (returns 0/1 tensor, no gradient)a.lt(b)a.gt(b)a.eq(b)a.isClose(b,tol)// Layouta.view(2,3)// reshapea.permute(1,0)// transposea.contiguous()// ensure contiguous memory// Linear algebraa.matmul(b)// matrix multiplication// Convolutiona.conv1d(weight,stride,padding)a.conv2d(weight,stride,padding)// Utilitiesa.clone()a.detach()// copy / detach from grapha.toString()// debug stringa.backward()// run backward passa.setRequiresGrad(true)// enable gradient tracking

Neural Network Modules

import{Linear,Conv1d,Conv2d,ReLU,Sigmoid,Embedding}from'@mni-ml/framework';constlinear=newLinear(inputSize,outputSize);constconv1d=newConv1d(inChannels,outChannels,kernelSize,stride,padding);constconv2d=newConv2d(inChannels,outChannels,kernelSize,stride,padding);constembedding=newEmbedding(vocabSize,embeddingDim);// Use in forward passconstout=linear.forward(input);

Functional Operations

import{softmax,gelu,layerNorm,crossEntropyLoss,dropout,avgpool2d,maxpool2d,tile}from'@mni-ml/framework';constsm=softmax(logits,dim);constg=gelu(x);constln=layerNorm(x,gamma,beta,eps);constloss=crossEntropyLoss(logits,targets);constdropped=dropout(x,rate,training);constpooled=avgpool2d(x,kernelH,kernelW);constmaxPooled=maxpool2d(x,kernelH,kernelW);consttiled=tile(x,[2,1]);

Optimizers

import{Adam,SGD}from'@mni-ml/framework';constoptimizer=newAdam(parameters,lr,beta1,beta2,eps,weightDecay);// orconstoptimizer=newSGD(parameters,lr);optimizer.step();// update parametersoptimizer.zeroGrad();// clear gradients

Backend Architecture

TypeScript API (tensor.ts, nn.ts, optimizer.ts)
│
└─→ N-API Bridge (lib.rs)
│
├─→ CPU Backend (Vec<f32>, 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.

# 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

License

MIT