TensorVM is a small neural network runtime built from scratch in Rust to understand and implement the systems underneath modern ML frameworks.
The runtime represents models as a tensor intermediate representation (IR), with the same computation graph used for forward execution, reverse-mode automatic differentiation, model serialization, and backend execution. I implemented the core tensor system, computation graph, autodiff engine, backend abstraction, portable model format, and Python API.
The core runtime is written in Rust and exposed to Python through PyO3, with a PyTorch-style frontend for defining models and optimizers. TensorVM currently supports CPU execution through a readable reference backend and an optimized backend, where blocked matrix multiplication delivers approximately 4.4x to 6.3x higher throughput on benchmarked matrix sizes.
I also built end-to-end MNIST training and serialization, allowing a trained model to be saved as a portable `.tvm` file, loaded without the original model code, and executed on a different backend while producing identical predictions.
The project is deliberately small and legible rather than production-scale. The goal was to understand how tensor runtimes, computational graphs, automatic differentiation, execution backends, and model serialization work underneath abstractions such as PyTorch and TVM.
Tech: Rust, Python, PyO3, NumPy, ndarray, automatic differentiation, tensor IR, compiler/runtime design, numerical computing, machine learning systems.
Open source : https://www.github.com/rayhankhilji/tensorvm