DEV Community

Shaurya Ojha
Shaurya Ojha

Posted on

I’m Building NumPy in Rust: Introducing rustnum

#rust #opensource #machinelearning #numpy

I’ve been working on rustnum, a pure-Rust, NumPy-style n-dimensional array library for Rust.

The goal is to make scientific computing and machine-learning-related numerical code more approachable for Rust developers while keeping the implementation portable and free from system BLAS or Fortran dependencies.

Why build another numerical array library?

Rust already has excellent numerical libraries, including ndarray. However, while working on ML infrastructure, I wanted to explore a different API and implementation strategy:

  • One dynamic-dimensional Array<T> type
  • NumPy-inspired function names
  • Familiar broadcasting and slicing
  • Pure-Rust numerical kernels
  • Optional parallelism through Rayon
  • Runtime SIMD acceleration where available

This project is also part of a larger experiment: rebuilding parts of the Python ML stack in Rust, one library at a time.

A NumPy-style Rust API 

use rustnum as np;

fn main() {
    let a = np::arange(0.0, 6.0, 1.0)
        .reshape(&[2, 3]);

    let b = np::array![10.0, 20.0, 30.0];

    // Broadcasting: [2, 3] + [3]
    let c = &a + &b;

    println!("{c}");
    println!("sum = {}", c.sum());
    println!("column means = {}", c.mean_axis(0));
}
Enter fullscreen mode Exit fullscreen mode

The API is intentionally familiar to anyone who has used NumPy:

a = np.arange(6.0).reshape(2, 3)
b = np.array([10.0, 20.0, 30.0])
c = a + b
Enter fullscreen mode Exit fullscreen mode

The Rust version uses compile-time element types such as Array<f64>, Array<i64>, and Array<bool>, while dimensions remain dynamic.

Current features

rustnum currently includes:

  • N-dimensional Array<T>
  • Broadcasting
  • Slicing with the s! macro
  • Copy-on-write array semantics
  • Element-wise arithmetic and comparisons
  • Reductions such as:
    • sum
    • mean
    • std
    • median
    • percentile
    • argmax
  • Shape manipulation:
    • reshape
    • transpose
    • concatenate
    • stack
    • split
    • flip
    • roll
    • sort
  • Linear algebra:
    • matmul
    • solve
    • inv
    • det
    • lstsq
    • cholesky
    • qr
    • eigh
    • svd
  • Random number generation
  • .npy and CSV file I/O
  • Optional Rayon parallelism
  • AVX2/FMA fast paths with portable scalar fallbacks

Packed GEMM in pure Rust

One of the more interesting parts of the project is matrix multiplication.

For f32 and f64, rustnum uses a packed GEMM implementation with cache blocking and AVX2/FMA micro-kernels when supported by the CPU.

The implementation has:

  • Packed matrix panels
  • Cache-aware blocking
  • AVX2/FMA kernels
  • Portable scalar fallback
  • Optional Rayon parallelism
  • Fused output transformations through matmul_map

There is no dependency on OpenBLAS, MKL, a C compiler, or a Fortran compiler.

A large matrix workload

I tested a configurable workload consisting of 12 iterations of two 8192 × 8192 matrix multiplications per iteration, applying tanh after each multiplication.

On one Windows 11 machine with 18 hardware threads:

Implementation Time
NumPy with OpenBLAS 237.686 seconds
rustnum 140.652 seconds

This is a single-machine reference measurement, not a universal performance claim. Results vary depending on the CPU, memory bandwidth, BLAS implementation, thread settings, and compiler version.

The workload can be run with:

$env:SIZE=8192
$env:ITERATIONS=12
cargo run --release -p rustnum --example heavy_matmul
Enter fullscreen mode Exit fullscreen mode

The equivalent Python workload is included in the repository as well.

Testing against NumPy

Correctness is more important than benchmark numbers for a numerical library.

The project includes:

  • Unit tests
  • Property-based tests
  • NumPy-generated parity fixtures
  • NPY round-trip tests
  • Random-number regression tests
  • Linear algebra tests
  • Tests with and without Rayon
  • Documentation tests

The main validation commands are:

cargo fmt --all
cargo clippy --all-targets -- -D warnings
cargo test --all-features
cargo test --no-default-features
cargo doc --no-deps
Enter fullscreen mode Exit fullscreen mode

Installation

The crate is available on crates.io:

cargo add rustnum
Enter fullscreen mode Exit fullscreen mode

Or add it manually:

[dependencies]
rustnum = "0.1.0"
Enter fullscreen mode Exit fullscreen mode

Documentation:

Current limitations

This is an early release and is not intended to replace mature numerical libraries yet.

Some functionality still planned or being developed includes:

  • einsum
  • More complete .npz support in the public API
  • More advanced indexing
  • Expanded FFT and polynomial APIs
  • Complex-number support
  • Python bindings
  • More optimized decompositions for SVD, QR, and LU
  • Broader platform and architecture testing

The API may also evolve based on real-world feedback.

What I’m looking for

I would especially appreciate feedback from Rust and scientific-computing developers about:

  1. Is a single dynamic-dimensional Array<T> type a useful API?
  2. How does this design compare with ndarray for your use cases?
  3. Which NumPy features should be prioritized next?
  4. Are the panic-versus-error boundaries intuitive?
  5. What numerical and performance benchmarks would you like to see?
  6. Would Python bindings make this project more useful to you?

Conclusion

rustnum started as an experiment in bringing NumPy-like APIs and numerical kernels to Rust. It has now grown into a working crate with broadcasting, linear algebra, random numbers, file I/O, SIMD kernels, and NumPy parity tests.

If you are interested in Rust-based scientific computing, I would love to hear your thoughts and suggestions.

Project: github.com/shauryaaojha/rustnum

Top comments (0)