mni-ml/framework: a TypeScript autograd API over Rust CPU, CUDA and WebGPU backends

The gap mni-ml/framework is trying to fill

Most people who want to learn how a tensor library works end up reading PyTorch source, which is C++ and CUDA spread across a very large repository. The other common route is a pure JavaScript teaching library, which is readable but stops at the CPU. This project sits between those two positions. It offers a TypeScript API that mirrors PyTorch naming closely enough to be familiar (Tensor, Module, Parameter, Adam, SGD, crossEntropyLoss) while the actual arithmetic lives in Rust, with three interchangeable backends behind it. The stated purpose in the repository description is to understand how ML frameworks and models work internally, and the architecture reflects that: the layer boundaries are explicit and the backend is a compile-time choice rather than a runtime plugin. The audience is therefore someone who wants to trace a backward pass from a TypeScript call into a Rust kernel and, if they choose, into a .cu or .wgsl file. It is not positioned as a faster alternative to a Python stack, and the README makes no throughput claims.

The N-API bridge and what the three backends share

The README gives the data flow as a diagram: TypeScript API files (tensor.ts, nn.ts, optimizer.ts) call into an N-API bridge in lib.rs, which dispatches to one of three backends. The CPU backend is pure Rust over Vec<f32>. The CUDA backend uses cudarc plus .cu kernel files. The WebGPU backend uses wgpu plus .wgsl shaders. Two details in that diagram matter more than the list itself. First, all three backends share the same autograd tape and tensor store, so switching backend does not change how gradients are recorded, only where the arithmetic runs. Second, the feature flags are mutually exclusive at compile time: cpu is the default and needs no GPU, cuda targets NVIDIA hardware, and webgpu reaches Metal, Vulkan and DX12 through wgpu. Mutually exclusive flags mean one binary carries one backend. You cannot build a single artifact that falls back from CUDA to CPU at runtime, and you cannot ship a package that picks a backend based on the user's machine without shipping several builds. That is a deliberate simplification for a learning project and a real constraint for anything else.

What the TypeScript surface actually exposes

The API reference lists a fairly conventional set. Tensor creation covers zeros, ones, rand, randn and fromFloat32. Arithmetic includes add, sub, mul, div, neg, exp, log and pow, each accepting either a tensor or a scalar where the README shows both forms. Activations are relu and sigmoid on the tensor, with gelu available as a functional operation. Reductions are sum, mean and max, each with a dim argument or without one for the whole tensor. Layout operations are view, permute and contiguous. Linear algebra is matmul. Convolution is exposed both as tensor methods (conv1d, conv2d with stride and padding) and as modules (Conv1d, Conv2d taking inChannels, outChannels, kernelSize, stride, padding). Functional operations include softmax, layerNorm, crossEntropyLoss, dropout, avgpool2d, maxpool2d and tile. Comparison operators (lt, gt, eq, isClose) return a 0/1 tensor and, as the README notes, carry no gradient. Optimizers are Adam with lr, beta1, beta2, eps and weightDecay, and SGD with lr. There is no mention of a DataLoader, a dataset abstraction, a serialization format for checkpoints, or mixed precision. Those absences shape what you can realistically build.

Getting a build running, and the toolchain you need first

The published entry point is npm install @mni-ml/framework. The README also documents building from source, which it says is only needed if you are contributing or want a custom build, and which requires Rust via rustup. The native build commands are npm run build:native for the default CPU backend, npm run build:native:cuda when the CUDA toolkit is present, and npm run build:native:webgpu for the wgpu path. TypeScript is compiled separately with npm run build. The quick start example is short enough to quote the shape of it: create tensors with Tensor.rand, instantiate Linear layers, run forward passes, call crossEntropyLoss, then loss.backward(), build a parameter list by spreading each layer's parameters(), construct Adam with a learning rate, and call step() followed by zeroGrad(). Note the ordering in that example. The optimizer is constructed after the backward pass in the snippet, which works for a single step but is not the usual pattern; in a training loop you would construct the optimizer once and call step() and zeroGrad() each iteration. The README does not show a full loop, so the loop structure is something you assemble yourself.

Where this framework stops being the right tool

The feature list has no checkpoint save or load, no device transfer API, no DataLoader and no batching utility, so a real training run means writing your own data pipeline in TypeScript and your own weight serialization. The backend flags being compile-time exclusive rules out the common deployment pattern where one published package adapts to whatever hardware the user has. The comparison operators returning non-differentiable tensors is standard, but combined with the absence of any documented masking or scatter operations it narrows what kinds of models you can express. There is also no release history in the material provided, and no version or stability statement in the README, so treat the API as moving. If you need to train a model that you intend to serve, or you need multi-GPU, or you need an ecosystem of pretrained weights, this is the wrong tool and the README does not pretend otherwise. The honest framing is that this is a framework you read and modify, not one you depend on.

How it differs from PyTorch and from pure-JS teaching libraries

Against PyTorch, the difference is not speed, it is the boundary. PyTorch gives you a Python front end over a C++ dispatcher with a very large operator set, autograd that handles arbitrary graphs including higher-order gradients, and a serialization format. mni-ml/framework gives you a TypeScript front end over an N-API bridge with a documented operator list that fits on one page, one autograd tape shared across backends, and no serialization. The interesting comparison is the WebGPU path: PyTorch does not target the browser, and running inference in a browser tab is a use case where a wgpu backend plus a TypeScript API is a coherent combination. Against pure-JS teaching libraries, the difference is that the heavy loops live in Rust and can be compiled to CUDA or WGSL, so the same TypeScript model code can run on a CPU Vec<f32> or on a GPU kernel. That portability across three backends from one API is the specific thing this project does that neither PyTorch nor a JavaScript-only library does.

Maintenance, licence and what to verify before you commit

The licence is MIT, which permits commercial use, modification and redistribution provided the copyright notice and permission notice are included. That is a permissive licence and imposes no copyleft obligation on your own code, but it also means no warranty and no support commitment from the authors. If you vendor the Rust core into a product, you are the maintainer of that fork. On upgrade cost, the material shows no releases retrieved, so there is no changelog to diff against and no version tags to pin. The practical consequence is that upgrading means reading the diff of tensor.ts, nn.ts and optimizer.ts yourself. The build path also carries cost: any contributor or CI runner that wants the CUDA or WebGPU backend needs the CUDA toolkit or a wgpu-capable environment, and the CPU default will silently be what you get if you run npm run build:native without the suffix. Verify the installed package's backend before you benchmark anything against it.

Editorial conclusion

Adopt mni-ml/framework if you want a small, readable TypeScript surface over a Rust tensor engine and you are willing to build the native addon yourself for CUDA or WebGPU. Do not adopt it if you need a production training stack with a release history, distributed execution or a model zoo; nothing in the material suggests any of those exist. Before committing, check whether the published npm package ships a prebuilt binary for your platform or whether npm install triggers a Rust toolchain build, and confirm which backend flag the installed binary was compiled with, because cpu, cuda and webgpu are mutually exclusive at compile time.

Official sources

1. Issues
2. License: MIT
3. mni-ml/framework on GitHub
4. Project website
5. README