@mni-ml/framework: TypeScript ML with Rust Backends for CPU, CUDA, and WebGPU

What @mni-ml/framework Is Built to Teach

The repository description says the project is built to understand how ML frameworks and models work internally. That framing matters. This is not an attempt to replace PyTorch or TensorFlow for production use. It is a framework sized to be readable: a TypeScript API layer, an N-API bridge in Rust, and three backend implementations that share the same autograd tape and tensor store.

For a TypeScript developer learning how backpropagation works, or how a GPU dispatch layer fits between user code and CUDA kernels, the codebase is intentionally small enough to trace from end to end. The autograd system records operations on a tape and replays them in reverse during the backward pass. Each backend implements the same operations in a different compute context: plain Vec<f32> for CPU, CUDA kernels via cudarc for NVIDIA GPUs, and WGSL shaders via wgpu for WebGPU targets.

The PyTorch-like API is a deliberate design choice. Developers already familiar with PyTorch will recognize Tensor, Module, Parameter, and optimizer patterns immediately, which reduces the time spent learning the API and increases the time available to study the backend implementation.

Installing and Running the First Model

Installation from the npm registry is one command:

npm install @mni-ml/framework

The package ships prebuilt native binaries for common platforms as optional dependencies. The package.json lists separate binary packages for darwin-arm64, darwin-x64, linux-arm64-gnu, linux-x64-gnu, linux-x64-gnu-cuda, win32-x64-msvc, and their WebGPU variants. npm selects the correct one for the current platform automatically.

A quick start example from the README shows the shape of training with the framework:

import { Tensor, Linear, Adam, softmax, crossEntropyLoss } from '@mni-ml/framework';

const x = Tensor.rand([32, 10]);

const layer1 = new Linear(10, 64);
const layer2 = new Linear(64, 3);

let h = layer1.forward(x).relu();
let logits = layer2.forward(h);
let loss = crossEntropyLoss(logits, targets);

loss.backward();

const params = [...layer1.parameters(), ...layer2.parameters()];
const optimizer = new Adam(params, 0.001);
optimizer.step();

This creates a two-layer classifier, runs a forward pass through it, computes a cross-entropy loss, calls backward to compute gradients, and takes one Adam optimizer step. The structure follows the PyTorch training loop closely, which means existing PyTorch knowledge transfers to reading this code.

Tensor Creation and Arithmetic Operations

The Tensor class supports four creation methods from the README:

Tensor.zeros([2, 3])
Tensor.ones([2, 3])
Tensor.rand([2, 3])
Tensor.randn([2, 3])

Arithmetic operations on tensors all carry autograd: add, sub, mul, div, neg, exp, log, and pow. Each supports both tensor-tensor and tensor-scalar forms where documented. The autograd tape records each operation as it runs, and calling .backward() on the loss tensor walks the tape in reverse to compute gradients for all Parameter tensors in the graph.

The functional operations available include softmax, gelu, layerNorm, crossEntropyLoss, dropout, avgpool2d, maxpool2d, and tile. These are imported directly from the package and applied as functions rather than as methods on a tensor.

Built-in module classes cover the core layer types: Linear, Conv1d, Conv2d, Embedding, ReLU, Sigmoid, and Tanh. Two optimizer implementations are provided: SGD and Adam. The Adam optimizer accepts beta1, beta2, epsilon, and weight decay parameters. The step method updates parameters in place, and zeroGrad clears accumulated gradients before the next forward pass.

Backend Architecture: Three Compute Paths from One API

The package uses a single TypeScript API that routes computation to one of three Rust backends through an N-API bridge. The backend is selected at compile time using mutually exclusive feature flags.

The README shows the architecture as a diagram:

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. The CPU backend uses plain Rust vectors, making it portable to any platform without special hardware. The CUDA backend uses cudarc and CUDA kernel files for NVIDIA GPUs. The WebGPU backend uses wgpu with WGSL shaders, which targets Metal on macOS, Vulkan on Linux, and DirectX 12 on Windows.

The npm package ships prebuilt binaries, so most users do not need to compile Rust at all. The build-from-source path exists for contributors or for custom builds, and it requires Rust to be installed.

Building from Source When You Need a Custom Backend

The README describes building from source as only necessary for contributors or custom builds. It requires Rust, which can be installed via rustup.rs.

Three build commands correspond to the three backends:

npm run build:native

This builds the CPU backend. For CUDA:

npm run build:native:cuda

This requires the CUDA toolkit to be installed. For WebGPU:

npm run build:native:webgpu

After building the native extension, compile the TypeScript layer:

npm run build

The feature flags are mutually exclusive at compile time. A given binary supports exactly one compute backend. To switch backends, you rebuild with the corresponding flag. This simplifies the runtime: there is no dynamic backend selection at inference time.

The Node.js requirement in package.json is 22.18.0 or newer. Teams using an older Node.js version will need to upgrade before installing.

Where This Package Falls Short

The last push to the repository was on 2026-04-20. There are no GitHub releases, no changelog, and no tagged versions. That means there is no stable release to pin against in a production dependency file. Any project that depends on @mni-ml/framework at a specific state must pin to a commit rather than a version number.

The README explicitly describes the project as built to understand how ML frameworks work internally. That is a statement of intent that sets expectations about scope. Comprehensive operator coverage, performance at production scale, distributed training, quantization, and model export are not goals documented anywhere in the repository.

The coverage of optimizers is limited to SGD and Adam. There is no AdaGrad, RMSProp, or learning rate scheduler implementation documented in the README. The Adam implementation accepts standard parameters including beta1, beta2, epsilon, and weight decay, but a more advanced training loop with a cosine annealing or step decay schedule would require the developer to implement the scheduler themselves.

There is no documented path for loading pretrained model weights from common formats like safetensors or ONNX. Building a project that uses pretrained weights with this framework would require writing custom loading code. The examples directory contains an xor.js file, and the toy directory holds additional small experiments, but neither directory provides guidance on more complex use cases like transfer learning or fine-tuning.

The Node.js requirement of version 22.18.0 or newer is a potential constraint for teams on older long-term-support Node.js versions that have not yet upgraded.

Comparison with PyTorch for Node.js Work

PyTorch does not have a native Node.js API. Developers who want to run PyTorch models in a Node.js server typically call into a Python subprocess, use the ONNX Runtime Node.js bindings, or deploy a separate model serving endpoint. @mni-ml/framework takes a different path: it is a native npm package that runs the full training and inference loop inside the Node.js process.

The difference in scope is large. PyTorch supports thousands of operators, distributed training, mixed precision, model quantization, and an ecosystem of tools for deployment. @mni-ml/framework documents the operators and modules in its README and makes no claims beyond them. For a developer who needs to train a small classifier inside a Node.js application and does not want to manage a Python subprocess, @mni-ml/framework is a more self-contained option. For any project that requires a specific operator not in the documented set, the missing coverage is a hard blocker.

The WebGPU backend is a meaningful distinction from PyTorch. PyTorch does not have a production WebGPU backend, so @mni-ml/framework is one of few options for running GPU-accelerated ML in contexts where WebGPU is the only available GPU interface, such as certain server configurations or future browser-adjacent runtimes.

Editorial conclusion

This package is a fit for TypeScript developers who want to train small neural networks without leaving the Node.js ecosystem, or who want to study how autograd and GPU dispatch work by reading a codebase that is small enough to follow from end to end. It is not a fit for production ML workloads: the last push was 2026-04-20, there are no GitHub releases to pin against, and the project's own README describes it as built to understand how ML frameworks work internally rather than to replace established frameworks. Before using it, verify that your Node.js version meets the stated requirement of 22.18.0 or newer, and confirm whether you need the CUDA or WebGPU backend by checking which native build command to run.

Frequently asked questions

Official sources

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